fix: track completed_at for cache eviction and resume

Author: DaleSeo

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

@@ -1,7 +1,7 @@
 use std::{
     collections::{HashMap, HashSet, VecDeque},
     num::ParseIntError,
-    time::Duration,
+    time::{Duration, Instant},
 };
 
 use futures::{Stream, StreamExt};
@@ -285,6 +285,7 @@ impl CachedTx {
 struct HttpRequestWise {
     resources: HashSet<ResourceKey>,
     tx: CachedTx,
+    completed_at: Option<Instant>,
 }
 
 type HttpRequestId = u64;
@@ -355,28 +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.get_mut(&http_request_id) {
-                        for resource in channel.resources.drain() {
-                            self.resource_router.remove(&resource);
-                        }
-                        // Replace the sender with a closed dummy so no new
-                        // messages are routed here, but the cache stays alive
-                        // for late resume requests.
-                        let (closed_tx, _) = tokio::sync::mpsc::channel(1);
-                        channel.tx.tx = closed_tx;
-                    }
-                }
-            } 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) {
@@ -413,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);
@@ -421,7 +426,6 @@ impl LocalSessionWorker {
     async fn establish_request_wise_channel(
         &mut self,
     ) -> Result<StreamableHttpMessageReceiver, SessionError> {
-        self.tx_router.retain(|_, rw| !rw.tx.tx.is_closed());
         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());
@@ -430,6 +434,7 @@ impl LocalSessionWorker {
             HttpRequestWise {
                 resources: Default::default(),
                 tx: CachedTx::new(tx, Some(http_request_id), starting_index),
+                completed_at: None,
             },
         );
         tracing::debug!(http_request_id, "establish new request wise channel");
@@ -540,29 +545,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) {
-                    let was_completed = request_wise.tx.tx.is_closed();
-                    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 was_completed {
-                        // Close 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,
-                    })
-                } else {
-                    tracing::debug!(
-                        http_request_id,
-                        "Request-wise channel not found, 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,
         }
@@ -972,6 +973,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() => {
@@ -1098,12 +1100,17 @@ pub struct SessionConfig {
     /// 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 {
@@ -1112,6 +1119,7 @@ impl Default for SessionConfig {
             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,46 @@ 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) => {
+                    // The referenced stream is gone (completed + evicted or
+                    // never existed). Fall through to create a fresh standalone
+                    // stream so EventSource auto-reconnection stays alive
+                    // without replaying events from a different stream.
+                    tracing::debug!("Resume failed ({e}), creating standalone stream");
+                }
+            }
         }
+        // Create standalone stream (also the fallback for failed resume)
+        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>

crates/rmcp/tests/test_streamable_http_priming.rs

@@ -197,115 +197,6 @@ async fn test_request_wise_priming_includes_http_request_id() -> anyhow::Result<
     Ok(())
 }
 
-#[tokio::test]
-async fn test_resume_after_request_wise_channel_completed() -> anyhow::Result<()> {
-    let ct = CancellationToken::new();
-
-    let service: StreamableHttpService<Calculator, LocalSessionManager> =
-        StreamableHttpService::new(
-            || Ok(Calculator::new()),
-            Default::default(),
-            StreamableHttpServerConfig::default()
-                .with_sse_keep_alive(None)
-                .with_cancellation_token(ct.child_token()),
-        );
-
-    let router = axum::Router::new().nest_service("/mcp", service);
-    let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
-    let addr = tcp_listener.local_addr()?;
-
-    let handle = tokio::spawn({
-        let ct = ct.clone();
-        async move {
-            let _ = axum::serve(tcp_listener, router)
-                .with_graceful_shutdown(async move { ct.cancelled_owned().await })
-                .await;
-        }
-    });
-
-    let client = reqwest::Client::new();
-
-    // Initialize session
-    let response = client
-        .post(format!("http://{addr}/mcp"))
-        .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?;
-    assert_eq!(response.status(), 200);
-    let session_id: SessionId = response.headers()["mcp-session-id"].to_str()?.into();
-
-    // Complete handshake
-    let status = client
-        .post(format!("http://{addr}/mcp"))
-        .header("Content-Type", "application/json")
-        .header("Accept", "application/json, text/event-stream")
-        .header("mcp-session-id", session_id.to_string())
-        .header("Mcp-Protocol-Version", "2025-06-18")
-        .body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
-        .send()
-        .await?
-        .status();
-    assert_eq!(status, 202);
-
-    // Call a tool and consume the full response (channel completes)
-    let body = client
-        .post(format!("http://{addr}/mcp"))
-        .header("Content-Type", "application/json")
-        .header("Accept", "application/json, text/event-stream")
-        .header("mcp-session-id", session_id.to_string())
-        .header("Mcp-Protocol-Version", "2025-06-18")
-        .body(r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"sum","arguments":{"a":1,"b":2}}}"#)
-        .send()
-        .await?
-        .text()
-        .await?;
-
-    let events: Vec<&str> = body.split("\n\n").filter(|e| !e.is_empty()).collect();
-    assert!(
-        events.len() >= 2,
-        "expected priming + response, got: {body}"
-    );
-    assert!(events[0].contains("id: 0/0"));
-    assert!(events[1].contains(r#""id":2"#));
-
-    // Resume with Last-Event-ID after the channel has completed.
-    // The cached events should be replayed and the stream should end.
-    let resume_response = client
-        .get(format!("http://{addr}/mcp"))
-        .header("Accept", "text/event-stream")
-        .header("mcp-session-id", session_id.to_string())
-        .header("Mcp-Protocol-Version", "2025-06-18")
-        .header("last-event-id", "0/0")
-        .timeout(std::time::Duration::from_secs(5))
-        .send()
-        .await?;
-    assert_eq!(resume_response.status(), 200);
-
-    let resume_body = resume_response.text().await?;
-    let resume_events: Vec<&str> = resume_body
-        .split("\n\n")
-        .filter(|e| !e.is_empty())
-        .collect();
-    assert!(
-        !resume_events.is_empty(),
-        "expected replayed events on resume, got empty"
-    );
-
-    // The replayed event should contain the original response
-    let replayed = resume_events[0];
-    assert!(
-        replayed.contains(r#""id":2"#),
-        "replayed event should contain the tool response, got: {replayed}"
-    );
-
-    ct.cancel();
-    handle.await?;
-
-    Ok(())
-}
-
 #[tokio::test]
 async fn test_priming_on_stream_close() -> anyhow::Result<()> {
     use std::sync::Arc;