From 14976da2a6090c195d6b12288c8d197b31703f7a Mon Sep 17 00:00:00 2001 From: freedom-winds Date: Tue, 11 Aug 2026 20:03:21 +0800 Subject: [PATCH 1/3] fix: preserve SSE resume cache and correct event replay --- .../streamable_http_server/session/local.rs | 42 +++++++++++++++++-- 1 file changed, 39 insertions(+), 3 deletions(-) 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 cc9e14893..1986d3b3d 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/local.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/local.rs @@ -335,7 +335,14 @@ impl CachedTx { .as_deref() .unwrap_or_default() .parse::()?; - let sync_index = index.saturating_sub(front_event_id.index); + if index < front_event_id.index { + return Err(SessionError::InvalidEventId); + } + // Last-Event-ID is the last event the client received, so resume only + // with events that follow it. + let sync_index = index + .saturating_sub(front_event_id.index) + .saturating_add(1); if sync_index > self.cache.len() { // invalid index return Err(SessionError::InvalidEventId); @@ -638,8 +645,14 @@ impl LocalSessionWorker { OutboundChannel::RequestWise { id, close } => { if let Some(request_wise) = self.tx_router.get_mut(&id) { request_wise.tx.send(message).await?; - if close && let Some(channel) = self.tx_router.remove(&id) { - for resource in channel.resources { + if close { + let resources: Vec<_> = request_wise.resources.drain().collect(); + request_wise.completed_at = Some(Instant::now()); + // Retain the completed channel until TTL eviction so a + // disconnected client can resume its final response. + let (closed_tx, _) = tokio::sync::mpsc::channel(1); + request_wise.tx.tx = closed_tx; + for resource in resources { self.resource_router.remove(&resource); } } @@ -1371,4 +1384,27 @@ mod sep2260_routing_tests { let channel = worker.resolve_outbound_channel(&roots_request(Some(originating_id))); assert!(matches!(channel, OutboundChannel::Common)); } + + #[tokio::test] + async fn completed_request_channel_retains_cache_for_resume() { + let (_handle, mut worker) = create_local_session("test-session", SessionConfig::default()); + let receiver = worker.establish_request_wise_channel().await.unwrap(); + let http_request_id = receiver.http_request_id.unwrap(); + let request_id = RequestId::Number(7); + worker.register_resource( + ResourceKey::McpRequestId(request_id.clone()), + http_request_id, + ); + + worker + .handle_server_message(ServerJsonRpcMessage::error( + crate::model::ErrorData::internal_error("failed", None), + Some(request_id), + )) + .await + .unwrap(); + + assert!(worker.tx_router[&http_request_id].completed_at.is_some()); + assert!(worker.tx_router[&http_request_id].resources.is_empty()); + } } From 4d8d06ac317c9b7e1c795020a55740fe1947f3c9 Mon Sep 17 00:00:00 2001 From: ZT Winds Date: Wed, 12 Aug 2026 13:37:17 +0800 Subject: [PATCH 2/3] fix: address SSE resume cache review feedback --- .../streamable_http_server/session/local.rs | 78 +++++++++++++++++-- 1 file changed, 72 insertions(+), 6 deletions(-) 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 1986d3b3d..ac2416782 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/local.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/local.rs @@ -335,14 +335,13 @@ impl CachedTx { .as_deref() .unwrap_or_default() .parse::()?; - if index < front_event_id.index { + // Last-Event-ID is the last event the client received. The resume is + // valid as long as the next event has not already fallen out of cache. + let next_index = index.saturating_add(1); + if next_index < front_event_id.index { return Err(SessionError::InvalidEventId); } - // Last-Event-ID is the last event the client received, so resume only - // with events that follow it. - let sync_index = index - .saturating_sub(front_event_id.index) - .saturating_add(1); + let sync_index = next_index.saturating_sub(front_event_id.index); if sync_index > self.cache.len() { // invalid index return Err(SessionError::InvalidEventId); @@ -668,6 +667,9 @@ impl LocalSessionWorker { &mut self, last_event_id: EventId, ) -> Result { + // A resume request can be the first event after a completed entry's + // TTL expires, so enforce eviction here before looking up its cache. + self.evict_expired_channels(); // Clean up closed shadow senders before processing self.shadow_txs.retain(|tx| !tx.is_closed()); @@ -1407,4 +1409,68 @@ mod sep2260_routing_tests { assert!(worker.tx_router[&http_request_id].completed_at.is_some()); assert!(worker.tx_router[&http_request_id].resources.is_empty()); } + + #[tokio::test] + async fn resume_accepts_event_immediately_before_cache_front() { + let (tx, mut rx) = tokio::sync::mpsc::channel(4); + let mut cached_tx = CachedTx::new(tx, None, 0, "test-stream".into(), None); + cached_tx.cache.push_back(ServerSseMessage { + event_id: Some("5".into()), + ..Default::default() + }); + cached_tx.cache.push_back(ServerSseMessage { + event_id: Some("6".into()), + ..Default::default() + }); + + assert!(matches!( + cached_tx.sync(3).await, + Err(SessionError::InvalidEventId) + )); + cached_tx.sync(4).await.unwrap(); + + assert_eq!(rx.recv().await.unwrap().event_id.as_deref(), Some("5")); + assert_eq!(rx.recv().await.unwrap().event_id.as_deref(), Some("6")); + } + + #[tokio::test] + async fn resume_rejects_expired_completed_request_cache() { + let config = SessionConfig { + completed_cache_ttl: Duration::ZERO, + ..SessionConfig::default() + }; + let (_handle, mut worker) = create_local_session("test-session", config); + let receiver = worker.establish_request_wise_channel().await.unwrap(); + let http_request_id = receiver.http_request_id.unwrap(); + let request_id = RequestId::Number(7); + worker.register_resource( + ResourceKey::McpRequestId(request_id.clone()), + http_request_id, + ); + worker + .handle_server_message(ServerJsonRpcMessage::error( + crate::model::ErrorData::internal_error("failed", None), + Some(request_id), + )) + .await + .unwrap(); + + let last_event_id = worker.tx_router[&http_request_id] + .tx + .cache + .front() + .unwrap() + .event_id + .as_deref() + .unwrap() + .parse::() + .unwrap(); + let result = worker.resume(last_event_id).await; + + assert!(matches!( + result, + Err(SessionError::ChannelClosed(Some(id))) if id == http_request_id + )); + assert!(!worker.tx_router.contains_key(&http_request_id)); + } } From 825fece75b7cb8af2d8e4ae92236d22d9ea1d00b Mon Sep 17 00:00:00 2001 From: ZT Winds Date: Wed, 12 Aug 2026 19:07:55 +0800 Subject: [PATCH 3/3] fix: preserve cache replay for fresh common streams --- .../streamable_http_server/session/local.rs | 26 ++++++++++++++----- .../rmcp/tests/test_sse_concurrent_streams.rs | 12 ++++----- 2 files changed, 25 insertions(+), 13 deletions(-) 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 ac2416782..1a361d1b5 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/local.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/local.rs @@ -355,6 +355,17 @@ impl CachedTx { } Ok(()) } + + async fn replay_cached(&mut self) -> Result<(), SessionError> { + for message in &self.cache { + let send_result = self.tx.send(message.clone()).await; + if send_result.is_err() { + let event_id: EventId = message.event_id.as_deref().unwrap_or_default().parse()?; + return Err(SessionError::ChannelClosed(Some(event_id.index as u64))); + } + } + Ok(()) + } } struct HttpRequestWise { @@ -705,17 +716,16 @@ impl LocalSessionWorker { 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 + self.resume_or_shadow_common(None).await } /// Resume the common channel, or create a shadow stream if the primary is /// still active. /// /// When the primary common channel is dead (receiver dropped), replace it - /// so this stream becomes the new primary notification channel. Cached - /// messages are replayed from `last_event_index` so the client receives - /// any events it missed (including server-initiated requests). + /// so this stream becomes the new primary notification channel. Resume + /// requests replay events after `last_event_index`; fresh GETs without a + /// Last-Event-ID replay the locally retained cache. /// /// When the primary is still active, create a "shadow" stream — an idle SSE /// connection kept alive by keep-alive pings. This prevents multiple @@ -736,8 +746,10 @@ impl LocalSessionWorker { // Primary common channel is dead — replace it. tracing::debug!("Replacing dead common channel with new primary"); self.common.tx = tx; - if let Some(last_event_index) = last_event_index { - self.common.sync(last_event_index).await?; + match last_event_index { + Some(last_event_index) => self.common.sync(last_event_index).await?, + None if self.event_store.is_none() => self.common.replay_cached().await?, + None => {} } } else { // Primary common channel is still active. Create a shadow stream diff --git a/crates/rmcp/tests/test_sse_concurrent_streams.rs b/crates/rmcp/tests/test_sse_concurrent_streams.rs index e1e885282..d41f171be 100644 --- a/crates/rmcp/tests/test_sse_concurrent_streams.rs +++ b/crates/rmcp/tests/test_sse_concurrent_streams.rs @@ -695,8 +695,8 @@ async fn dropping_shadows_does_not_affect_primary() { // ─── Tests: Cache replay on dead primary replacement ───────────────────────── /// When a notification is sent while the primary is alive, then the primary -/// dies and a new GET resumes with Last-Event-ID "0", the replacement primary -/// should receive the cached notification via sync() replay. +/// dies and a fresh GET connects without Last-Event-ID, the replacement primary +/// should receive the locally cached notification. #[tokio::test] async fn dead_primary_replacement_replays_cached_events() { let ct = CancellationToken::new(); @@ -721,15 +721,15 @@ async fn dead_primary_replacement_replays_cached_events() { drop(get1); tokio::time::sleep(Duration::from_millis(100)).await; - // Resume with Last-Event-ID "0" — primary is dead, should replace it - // and replay cached events from index 0 - let get_resume = open_resume_get(&client, &url, &session_id, "0").await; + // Fresh GET without Last-Event-ID — primary is dead, so it should replace + // the old primary and replay the locally retained cache. + let get_resume = open_standalone_get(&client, &url, &session_id).await; assert_eq!(get_resume.status(), 200); // The cached notification should be replayed on the new primary assert!( wait_for_sse_event(get_resume, "tools/list_changed", Duration::from_secs(3)).await, - "Replacement primary should receive cached notification via sync() replay" + "Replacement primary should receive cached notification" ); ct.cancel();