diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 280cfd376..757f69056 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -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 = [ diff --git a/crates/rmcp/src/transport/common/auth/streamable_http_client.rs b/crates/rmcp/src/transport/common/auth/streamable_http_client.rs index f0a6211af..2069e6d31 100644 --- a/crates/rmcp/src/transport/common/auth/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/auth/streamable_http_client.rs @@ -31,7 +31,7 @@ where async fn get_stream( &self, uri: std::sync::Arc, - session_id: std::sync::Arc, + session_id: Option>, last_event_id: Option, mut auth_token: Option, custom_headers: HashMap, @@ -50,7 +50,7 @@ where async fn get_stream_with_max_sse_event_size( &self, uri: std::sync::Arc, - session_id: std::sync::Arc, + session_id: Option>, last_event_id: Option, mut auth_token: Option, custom_headers: HashMap, diff --git a/crates/rmcp/src/transport/common/client_side_sse.rs b/crates/rmcp/src/transport/common/client_side_sse.rs index ba21657f7..ce425c3a2 100644 --- a/crates/rmcp/src/transport/common/client_side_sse.rs +++ b/crates/rmcp/src/transport/common/client_side_sse.rs @@ -288,6 +288,7 @@ pin_project_lite::pin_project! { where R: SseStreamReconnect { retry_policy: Arc, + reconnect_only_after_event_id: bool, last_event_id: Option, server_retry_interval: Option, connector: R, @@ -304,6 +305,22 @@ impl SseAutoReconnectStream { ) -> 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, + ) -> Self { + Self { + retry_policy, + reconnect_only_after_event_id: true, last_event_id: None, server_retry_interval: None, connector, @@ -317,6 +334,7 @@ impl SseAutoReconnectStream> { 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); + } } diff --git a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs index 7032e1a87..d2557dc0e 100644 --- a/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs +++ b/crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs @@ -52,7 +52,7 @@ impl StreamableHttpClient for reqwest::Client { async fn get_stream( &self, uri: Arc, - session_id: Arc, + session_id: Option>, last_event_id: Option, auth_token: Option, custom_headers: HashMap, @@ -71,7 +71,7 @@ impl StreamableHttpClient for reqwest::Client { async fn get_stream_with_max_sse_event_size( &self, uri: Arc, - session_id: Arc, + session_id: Option>, last_event_id: Option, auth_token: Option, custom_headers: HashMap, @@ -79,8 +79,10 @@ impl StreamableHttpClient for reqwest::Client { ) -> Result>, StreamableHttpError> { 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); } diff --git a/crates/rmcp/src/transport/common/server_side_http.rs b/crates/rmcp/src/transport/common/server_side_http.rs index 4969ff793..09df6d285 100644 --- a/crates/rmcp/src/transport/common/server_side_http.rs +++ b/crates/rmcp/src/transport/common/server_side_http.rs @@ -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( diff --git a/crates/rmcp/src/transport/common/unix_socket.rs b/crates/rmcp/src/transport/common/unix_socket.rs index 899548313..ef6555b7f 100644 --- a/crates/rmcp/src/transport/common/unix_socket.rs +++ b/crates/rmcp/src/transport/common/unix_socket.rs @@ -376,7 +376,7 @@ impl StreamableHttpClient for UnixSocketHttpClient { async fn get_stream( &self, uri: Arc, - session_id: Arc, + session_id: Option>, last_event_id: Option, auth_token: Option, custom_headers: HashMap, @@ -396,7 +396,7 @@ impl StreamableHttpClient for UnixSocketHttpClient { async fn get_stream_with_max_sse_event_size( &self, uri: Arc, - session_id: Arc, + session_id: Option>, last_event_id: Option, auth_token: Option, custom_headers: HashMap, @@ -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); diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index 432ff705c..27e09c653 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -348,10 +348,14 @@ pub trait StreamableHttpClient: Clone + Send + 'static { auth_header: Option, custom_headers: HashMap, ) -> impl Future>> + 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, - session_id: Arc, + session_id: Option>, last_event_id: Option, auth_header: Option, custom_headers: HashMap, @@ -375,7 +379,7 @@ pub trait StreamableHttpClient: Clone + Send + 'static { fn get_stream_with_max_sse_event_size( &self, uri: Arc, - session_id: Arc, + session_id: Option>, last_event_id: Option, auth_header: Option, custom_headers: HashMap, @@ -399,7 +403,7 @@ pub struct RetryConfig { struct StreamableHttpClientReconnect { pub client: C, - pub session_id: Arc, + pub session_id: Option>, pub uri: Arc, pub auth_header: Option, pub custom_headers: HashMap, @@ -550,37 +554,6 @@ impl StreamableHttpClientWorker { Ok(()) } - /// Convert a raw SSE stream into a JSON-RPC message stream without - /// reconnection logic. - fn raw_sse_to_jsonrpc( - stream: BoxedSseStream, - ) -> impl Stream>> + 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::(&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 StreamableHttpClientWorker { fn reconnecting_sse_to_jsonrpc( stream: BoxedSseStream, client: C, - session_id: Arc, + session_id: Option>, uri: Arc, auth_header: Option, custom_headers: HashMap, @@ -598,7 +571,7 @@ impl StreamableHttpClientWorker { retry_config: Arc, ) -> impl Stream>> + Send + 'static { - SseAutoReconnectStream::new( + SseAutoReconnectStream::new_after_event_id( stream, StreamableHttpClientReconnect { client, @@ -614,10 +587,8 @@ impl StreamableHttpClientWorker { /// 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>, @@ -628,20 +599,17 @@ impl StreamableHttpClientWorker { max_sse_event_size: usize, retry_config: Arc, ) -> BoxStream<'static, Result>> { - 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 StreamableHttpClientWorker { 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 StreamableHttpClientWorker { 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 Worker for StreamableHttpClientWorker { /// async fn get_stream( /// &self, /// _uri: Arc, -/// _session_id: Arc, +/// _session_id: Option>, /// _last_event_id: Option, /// _auth_header: Option, /// _custom_headers: HashMap, @@ -1646,7 +1614,7 @@ impl StreamableHttpClientTransport { /// async fn get_stream( /// &self, /// _uri: Arc, - /// _session_id: Arc, + /// _session_id: Option>, /// _last_event_id: Option, /// _auth_header: Option, /// _custom_headers: HashMap, @@ -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, Option); + + #[derive(Clone, Default)] + struct StatelessReconnectClient { + reconnects: Arc>>, + } + + impl StreamableHttpClient for StatelessReconnectClient { + type Error = std::io::Error; + + async fn post_message( + &self, + _uri: Arc, + _message: ClientJsonRpcMessage, + _session_id: Option>, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result> { + Err(StreamableHttpError::UnexpectedServerResponse( + "unexpected POST".into(), + )) + } + + async fn delete_session( + &self, + _uri: Arc, + _session_id: Arc, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result<(), StreamableHttpError> { + Ok(()) + } + + async fn get_stream( + &self, + _uri: Arc, + session_id: Option>, + last_event_id: Option, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result> { + 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::::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", diff --git a/crates/rmcp/src/transport/streamable_http_server/session.rs b/crates/rmcp/src/transport/streamable_http_server/session.rs index ab0ff3244..28bd98bc2 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session.rs @@ -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, Self::Error>> + Send { futures::future::ready(Ok(RestoreOutcome::NotSupported)) } + + /// Return the shared event store used for resumable SSE streams. + fn event_store(&self) -> Option> { + None + } } diff --git a/crates/rmcp/src/transport/streamable_http_server/session/local.rs b/crates/rmcp/src/transport/streamable_http_server/session/local.rs index 7e2893206..ca88c088a 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/local.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/local.rs @@ -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>, pub session_config: SessionConfig, + event_store: Option>, +} + +impl LocalSessionManager { + /// Configure this session manager to use a shared event store. + pub fn with_event_store(mut self, event_store: Arc) -> Self { + self.event_store = Some(event_store); + self + } } #[derive(Debug, Error)] @@ -49,7 +59,11 @@ impl SessionManager for LocalSessionManager { type Transport = WorkerTransport; 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 + 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> { + self.event_store.clone() + } } /// `/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, @@ -217,6 +240,8 @@ struct CachedTx { http_request_id: Option, capacity: usize, starting_index: usize, + stream_id: StreamId, + event_store: Option>, } impl CachedTx { @@ -224,6 +249,8 @@ impl CachedTx { tx: Sender, http_request_id: Option, starting_index: usize, + stream_id: StreamId, + event_store: Option>, ) -> 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) -> Self { - Self::new(tx, None, 0) + fn new_common( + tx: Sender, + session_id: &SessionId, + event_store: Option>, + ) -> 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) -> Result<(), SessionError> { + self.store_cache_and_send(ServerSseMessage::retry(retry)) + .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 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>, event_rx: Receiver, session_config: SessionConfig, + event_store: Option>, } 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 for std::io::Error { @@ -450,12 +501,21 @@ impl LocalSessionWorker { ) -> Result { 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 { + 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, ) -> Result { 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, responder: oneshot::Sender>, }, + EstablishCommonChannel { + responder: oneshot::Sender>, + }, } #[derive(Debug, Clone)] @@ -820,13 +893,7 @@ impl LocalSessionHandle { ) -> Result { 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, config: SessionConfig, +) -> (LocalSessionHandle, LocalSessionWorker) { + create_local_session_with_event_store(id, config, None) +} + +fn create_local_session_with_event_store( + id: impl Into, + config: SessionConfig, + event_store: Option>, ) -> (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) } diff --git a/crates/rmcp/src/transport/streamable_http_server/session/never.rs b/crates/rmcp/src/transport/streamable_http_server/session/never.rs index a2f72d820..d83031c37 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/never.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/never.rs @@ -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>, +} + +impl NeverSessionManager { + /// Configure resumable SSE storage without enabling sessions. + pub fn with_event_store(mut self, event_store: Arc) -> Self { + self.event_store = Some(event_store); + self + } +} #[non_exhaustive] pub enum NeverTransport {} impl Transport for NeverTransport { @@ -107,4 +119,8 @@ impl SessionManager for NeverSessionManager { ) -> impl Future> + Send { futures::future::ready(Err(ErrorSessionManagementNotSupported)) } + + fn event_store(&self) -> Option> { + self.event_store.clone() + } } diff --git a/crates/rmcp/src/transport/streamable_http_server/session/store.rs b/crates/rmcp/src/transport/streamable_http_server/session/store.rs index e9a6de2d8..8b656201c 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/store.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/store.rs @@ -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 + Send + Sync + 'static>>; + +/// Type alias for boxed event store errors. +pub type EventStoreError = Box; + +/// 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; + + /// 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; +} + +impl std::fmt::Debug for dyn EventStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("") + } +} /// State persisted to an external store for cross-instance session recovery. /// diff --git a/crates/rmcp/src/transport/streamable_http_server/tower.rs b/crates/rmcp/src/transport/streamable_http_server/tower.rs index 16fa1eccb..f66235ce8 100644 --- a/crates/rmcp/src/transport/streamable_http_server/tower.rs +++ b/crates/rmcp/src/transport/streamable_http_server/tower.rs @@ -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>, +) -> 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, message: impl Into>, @@ -920,6 +937,98 @@ where (self.service_factory)() } + fn persisted_stateless_stream( + &self, + first: Option, + mut receiver: tokio::sync::mpsc::Receiver, + request_ct: CancellationToken, + event_store: Arc, + ) -> ReceiverStream { + 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, + receiver: tokio::sync::mpsc::Receiver, + 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::::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 { #[pin] inner: S, diff --git a/crates/rmcp/tests/test_streamable_http_event_store.rs b/crates/rmcp/tests/test_streamable_http_event_store.rs new file mode 100644 index 000000000..cc88f1077 --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_event_store.rs @@ -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>>); + +#[async_trait::async_trait] +impl SessionStore for InMemorySessionStore { + async fn load(&self, session_id: &str) -> Result, 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>>, + next_id: Arc, +} + +#[async_trait::async_trait] +impl EventStore for InMemoryEventStore { + async fn store_event( + &self, + stream_id: &str, + event: &ServerSseMessage, + ) -> Result { + 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 { + 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::>(); + 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, + ) -> Result { + 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, + event_store: Arc, + 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 = Arc::new(InMemorySessionStore::default()); + let event_store: Arc = 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 = 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(()) +} diff --git a/crates/rmcp/tests/test_streamable_http_stale_session.rs b/crates/rmcp/tests/test_streamable_http_stale_session.rs index 137460f7b..be1ac0269 100644 --- a/crates/rmcp/tests/test_streamable_http_stale_session.rs +++ b/crates/rmcp/tests/test_streamable_http_stale_session.rs @@ -160,7 +160,7 @@ impl StreamableHttpClient for ReinitDropsAcceptedResponseClient { async fn get_stream( &self, _uri: Arc, - session_id: Arc, + session_id: Option>, _last_event_id: Option, _auth_header: Option, _custom_headers: HashMap, @@ -168,7 +168,7 @@ impl StreamableHttpClient for ReinitDropsAcceptedResponseClient { futures::stream::BoxStream<'static, Result>, StreamableHttpError, > { - 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;