test: add long-running tool resume tests

Author: DaleSeo

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

@@ -356,29 +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);
-                        }
-                        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;
-                    }
-                }
-            } 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) {

crates/rmcp/tests/test_streamable_http_priming.rs

@@ -270,34 +270,40 @@ async fn test_resume_after_request_wise_channel_completed() -> anyhow::Result<()
     assert!(events[0].contains("id: 0/0"));
     assert!(events[1].contains(r#""id":2"#));
 
+    // Allow the session worker to finish processing the response
+    // and mark the channel as completed before resuming.
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
     // Resume with Last-Event-ID after the channel has completed.
-    // The cached events should be replayed and the stream should end.
+    // The cached events should be replayed.
     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];
+    // Read incrementally — the resumed stream may stay open as a
+    // standalone fallback, so we read just enough to verify replay.
+    use futures::StreamExt;
+    let mut stream = resume_response.bytes_stream();
+    let mut buf = String::new();
+    let read_result = tokio::time::timeout(Duration::from_secs(10), async {
+        while let Some(Ok(chunk)) = stream.next().await {
+            buf.push_str(&String::from_utf8_lossy(&chunk));
+            if buf.contains(r#""id":2"#) {
+                return true;
+            }
+        }
+        false
+    })
+    .await;
     assert!(
-        replayed.contains(r#""id":2"#),
-        "replayed event should contain the tool response, got: {replayed}"
+        read_result.unwrap_or(false),
+        "expected replayed tool response on resume, got: {buf}"
     );
 
     ct.cancel();