fix: include http_request_id in request-wise priming event IDs (#799)

* fix: include http_request_id in request-wise priming event IDs

* refactor: use Option::into_iter and usize::from for priming

* fix: retain event cache for completed request-wise channels

* fix: track completed_at for cache eviction and resume

* fix: log resume failures at warn level

* test: add completed_cache_ttl eviction test

* fix: return empty stream on failed resume

* test: add resume after completion test

Author: DaleSeo

crates/rmcp/src/transport/streamable_http_server/session/local.rs

@@ -1,10 +1,10 @@
 use std::{
     collections::{HashMap, HashSet, VecDeque},
     num::ParseIntError,
-    time::Duration,
+    time::{Duration, Instant},
 };
 
-use futures::Stream;
+use futures::{Stream, StreamExt};
 use thiserror::Error;
 use tokio::sync::{
     mpsc::{Receiver, Sender},
@@ -86,10 +86,17 @@ impl SessionManager for LocalSessionManager {
             .get(id)
             .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
         let receiver = handle.establish_request_wise_channel().await?;
-        handle
-            .push_message(message, receiver.http_request_id)
-            .await?;
-        Ok(ReceiverStream::new(receiver.inner))
+        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)))
     }
 
     async fn create_standalone_stream(
@@ -188,23 +195,29 @@ struct CachedTx {
     cache: VecDeque<ServerSseMessage>,
     http_request_id: Option<HttpRequestId>,
     capacity: usize,
+    starting_index: usize,
 }
 
 impl CachedTx {
-    fn new(tx: Sender<ServerSseMessage>, http_request_id: Option<HttpRequestId>) -> Self {
+    fn new(
+        tx: Sender<ServerSseMessage>,
+        http_request_id: Option<HttpRequestId>,
+        starting_index: usize,
+    ) -> Self {
         Self {
             cache: VecDeque::with_capacity(tx.capacity()),
             capacity: tx.capacity(),
             tx,
             http_request_id,
+            starting_index,
         }
     }
     fn new_common(tx: Sender<ServerSseMessage>) -> Self {
-        Self::new(tx, None)
+        Self::new(tx, None, 0)
     }
 
     fn next_event_id(&self) -> EventId {
-        let index = self.cache.back().map_or(0, |m| {
+        let index = self.cache.back().map_or(self.starting_index, |m| {
             m.event_id
                 .as_deref()
                 .unwrap_or_default()
@@ -272,6 +285,7 @@ impl CachedTx {
 struct HttpRequestWise {
     resources: HashSet<ResourceKey>,
     tx: CachedTx,
+    completed_at: Option<Instant>,
 }
 
 type HttpRequestId = u64;
@@ -342,23 +356,27 @@ pub struct StreamableHttpMessageReceiver {
 
 impl LocalSessionWorker {
     fn unregister_resource(&mut self, resource: &ResourceKey) {
-        if let Some(http_request_id) = self.resource_router.remove(resource) {
-            tracing::trace!(?resource, http_request_id, "unregister resource");
-            if let Some(channel) = self.tx_router.get_mut(&http_request_id) {
-                // It's okey to do so, since we don't handle batch json rpc request anymore
-                // and this can be refactored after the batch request is removed in the coming version.
-                if channel.resources.is_empty() || matches!(resource, ResourceKey::McpRequestId(_))
-                {
-                    tracing::debug!(http_request_id, "close http request wise channel");
-                    if let Some(channel) = self.tx_router.remove(&http_request_id) {
-                        for resource in channel.resources {
-                            self.resource_router.remove(&resource);
-                        }
-                    }
-                }
-            } else {
-                tracing::warn!(http_request_id, "http request wise channel not found");
-            }
+        let Some(http_request_id) = self.resource_router.remove(resource) else {
+            return;
+        };
+        tracing::trace!(?resource, http_request_id, "unregister resource");
+        let Some(channel) = self.tx_router.get_mut(&http_request_id) else {
+            tracing::warn!(http_request_id, "http request wise channel not found");
+            return;
+        };
+        if !channel.resources.is_empty() && !matches!(resource, ResourceKey::McpRequestId(_)) {
+            return;
+        }
+        tracing::debug!(http_request_id, "close http request wise channel");
+        let resources: Vec<_> = channel.resources.drain().collect();
+        channel.completed_at = Some(Instant::now());
+        // Close the sender so the client's SSE stream ends,
+        // but keep the entry so the cache is available for
+        // late resume requests.
+        let (closed_tx, _) = tokio::sync::mpsc::channel(1);
+        channel.tx.tx = closed_tx;
+        for resource in resources {
+            self.resource_router.remove(&resource);
         }
     }
     fn register_resource(&mut self, resource: ResourceKey, http_request_id: HttpRequestId) {
@@ -395,6 +413,11 @@ impl LocalSessionWorker {
             self.unregister_resource(&resource);
         }
     }
+    fn evict_expired_channels(&mut self) {
+        let ttl = self.session_config.completed_cache_ttl;
+        self.tx_router
+            .retain(|_, rw| rw.completed_at.is_none_or(|at| at.elapsed() < ttl));
+    }
     fn next_http_request_id(&mut self) -> HttpRequestId {
         let id = self.next_http_request_id;
         self.next_http_request_id = self.next_http_request_id.wrapping_add(1);
@@ -405,11 +428,13 @@ 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());
         self.tx_router.insert(
             http_request_id,
             HttpRequestWise {
                 resources: Default::default(),
-                tx: CachedTx::new(tx, Some(http_request_id)),
+                tx: CachedTx::new(tx, Some(http_request_id), starting_index),
+                completed_at: None,
             },
         );
         tracing::debug!(http_request_id, "establish new request wise channel");
@@ -524,28 +549,25 @@ impl LocalSessionWorker {
 
         match last_event_id.http_request_id {
             Some(http_request_id) => {
-                if let Some(request_wise) = self.tx_router.get_mut(&http_request_id) {
-                    // Resume existing request-wise channel
-                    let channel = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
-                    let (tx, rx) = channel;
-                    request_wise.tx.tx = tx;
-                    let index = last_event_id.index;
-                    // sync messages after index
-                    request_wise.tx.sync(index).await?;
-                    Ok(StreamableHttpMessageReceiver {
-                        http_request_id: Some(http_request_id),
-                        inner: rx,
-                    })
-                } else {
-                    // Request-wise channel completed (POST response already delivered).
-                    // The client's EventSource is reconnecting after the POST SSE stream
-                    // ended. Fall through to common channel handling below.
-                    tracing::debug!(
-                        http_request_id,
-                        "Request-wise channel completed, falling back to common channel"
-                    );
-                    self.resume_or_shadow_common(last_event_id.index).await
+                let request_wise = self
+                    .tx_router
+                    .get_mut(&http_request_id)
+                    .ok_or(SessionError::ChannelClosed(Some(http_request_id)))?;
+                let is_completed = request_wise.completed_at.is_some();
+                let (tx, rx) = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
+                request_wise.tx.tx = tx;
+                let index = last_event_id.index;
+                request_wise.tx.sync(index).await?;
+                if is_completed {
+                    // Drop the sender after replaying so the stream ends
+                    // instead of hanging indefinitely.
+                    let (closed_tx, _) = tokio::sync::mpsc::channel(1);
+                    request_wise.tx.tx = closed_tx;
                 }
+                Ok(StreamableHttpMessageReceiver {
+                    http_request_id: Some(http_request_id),
+                    inner: rx,
+                })
             }
             None => self.resume_or_shadow_common(last_event_id.index).await,
         }
@@ -955,6 +977,7 @@ impl Worker for LocalSessionWorker {
         let ct = context.cancellation_token.clone();
         let keep_alive = self.session_config.keep_alive.unwrap_or(Duration::MAX);
         loop {
+            self.evict_expired_channels();
             let keep_alive_timeout = tokio::time::sleep(keep_alive);
             let event = tokio::select! {
                 event = self.event_rx.recv() => {
@@ -1076,18 +1099,31 @@ pub struct SessionConfig {
     /// Defaults to 5 minutes. Set to `None` to disable (not recommended
     /// for long-running servers behind proxies).
     pub keep_alive: Option<Duration>,
+    /// SSE retry interval for priming events on request-wise streams.
+    /// When set, the session layer prepends a priming event with the correct
+    /// stream-identifying event ID to each request-wise SSE stream.
+    /// Default is 3 seconds, matching `StreamableHttpServerConfig::default()`.
+    pub sse_retry: Option<Duration>,
+    /// How long to retain completed request-wise channel caches for late
+    /// resume requests. After this duration, completed entries are evicted
+    /// and resume will return an error. Default is 60 seconds.
+    pub completed_cache_ttl: Duration,
 }
 
 impl SessionConfig {
     pub const DEFAULT_CHANNEL_CAPACITY: usize = 16;
     pub const DEFAULT_KEEP_ALIVE: Duration = Duration::from_secs(300);
+    pub const DEFAULT_SSE_RETRY: Duration = Duration::from_secs(3);
+    pub const DEFAULT_COMPLETED_CACHE_TTL: Duration = Duration::from_secs(60);
 }
 
 impl Default for SessionConfig {
     fn default() -> Self {
         Self {
             channel_capacity: Self::DEFAULT_CHANNEL_CAPACITY,
             keep_alive: Some(Self::DEFAULT_KEEP_ALIVE),
+            sse_retry: Some(Self::DEFAULT_SSE_RETRY),
+            completed_cache_ttl: Self::DEFAULT_COMPLETED_CACHE_TTL,
         }
     }
 }

crates/rmcp/src/transport/streamable_http_server/tower.rs

@@ -478,40 +478,52 @@ where
             .and_then(|v| v.to_str().ok())
             .map(|s| s.to_owned());
         if let Some(last_event_id) = last_event_id {
-            // check if session has this event id
-            let stream = self
+            match self
                 .session_manager
                 .resume(&session_id, last_event_id)
                 .await
-                .map_err(internal_error_response("resume session"))?;
-            // Resume doesn't need priming - client already has the event ID
-            Ok(sse_stream_response(
-                stream,
-                self.config.sse_keep_alive,
-                self.config.cancellation_token.child_token(),
-            ))
-        } else {
-            // create standalone stream
-            let stream = self
-                .session_manager
-                .create_standalone_stream(&session_id)
-                .await
-                .map_err(internal_error_response("create standalone stream"))?;
-            // Prepend priming event if sse_retry configured
-            let stream = if let Some(retry) = self.config.sse_retry {
-                let priming = ServerSseMessage::priming("0", retry);
-                futures::stream::once(async move { priming })
-                    .chain(stream)
-                    .left_stream()
-            } else {
-                stream.right_stream()
-            };
-            Ok(sse_stream_response(
-                stream,
-                self.config.sse_keep_alive,
-                self.config.cancellation_token.child_token(),
-            ))
+            {
+                Ok(stream) => {
+                    return Ok(sse_stream_response(
+                        stream,
+                        self.config.sse_keep_alive,
+                        self.config.cancellation_token.child_token(),
+                    ));
+                }
+                Err(e) => {
+                    // Return 200 with an immediately-closed empty stream.
+                    // Returning an HTTP error would cause EventSource to retry
+                    // with the same Last-Event-ID in an infinite loop. An empty
+                    // 200 cleanly terminates the EventSource without delivering
+                    // events from a different stream.
+                    tracing::warn!("Resume failed ({e}), returning empty stream");
+                    return Ok(sse_stream_response(
+                        futures::stream::empty(),
+                        None,
+                        self.config.cancellation_token.child_token(),
+                    ));
+                }
+            }
         }
+        // No Last-Event-ID — create standalone stream
+        let stream = self
+            .session_manager
+            .create_standalone_stream(&session_id)
+            .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);
+            futures::stream::once(async move { priming })
+                .chain(stream)
+                .left_stream()
+        } else {
+            stream.right_stream()
+        };
+        Ok(sse_stream_response(
+            stream,
+            self.config.sse_keep_alive,
+            self.config.cancellation_token.child_token(),
+        ))
     }
 
     async fn handle_post<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
@@ -598,20 +610,14 @@ where
 
                 match message {
                     ClientJsonRpcMessage::Request(_) => {
+                        // Priming for request-wise streams is handled by the
+                        // session layer (SessionManager::create_stream) which
+                        // has access to the http_request_id for correct event IDs.
                         let stream = self
                             .session_manager
                             .create_stream(&session_id, message)
                             .await
                             .map_err(internal_error_response("get session"))?;
-                        // Prepend priming event if sse_retry configured
-                        let stream = if let Some(retry) = self.config.sse_retry {
-                            let priming = ServerSseMessage::priming("0", retry);
-                            futures::stream::once(async move { priming })
-                                .chain(stream)
-                                .left_stream()
-                        } else {
-                            stream.right_stream()
-                        };
                         Ok(sse_stream_response(
                             stream,
                             self.config.sse_keep_alive,