From 14fc2f39102e815f4b9423a539006048019db16c Mon Sep 17 00:00:00 2001 From: Amey Pawar <138877912+ameyypawar@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:56:00 +0530 Subject: [PATCH 1/3] fix(transport): cancel in-flight request on streamable-http client disconnect (#857) When a streamable-HTTP client closes its TCP connection while a tool handler is still running, the request-wise SSE response stream is dropped, but nothing fired the per-request cancellation token. Only a natural response completion or an explicit notifications/cancelled cancelled it, so a long-running handler ran to completion after the client had already disconnected. Wrap the request-wise response stream in a drop-guard. When the stream is dropped before completing (client disconnect), it injects a synthetic notifications/cancelled for the request, reusing the existing cancellation path so the service fires the handler's RequestContext::ct. Streams that end normally (response delivered, or a reconnection handoff) do not cancel. Adds a regression test covering both the disconnect-cancels and the normal-completion paths. No public API change. --- .../streamable_http_server/session/local.rs | 104 ++++++++- .../test_streamable_http_disconnect_cancel.rs | 214 ++++++++++++++++++ 2 files changed, 314 insertions(+), 4 deletions(-) create mode 100644 crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs 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 7e289320..300f36cf 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,8 @@ use std::{ collections::{HashMap, HashSet, VecDeque}, num::ParseIntError, + pin::Pin, + task::{Context, Poll}, time::{Duration, Instant}, }; @@ -16,9 +18,10 @@ use tracing::instrument; use crate::{ RoleServer, model::{ - CancelledNotificationParam, ClientJsonRpcMessage, ClientNotification, ClientRequest, - JsonRpcNotification, JsonRpcRequest, Notification, ProgressNotificationParam, - ProgressToken, RequestId, ServerJsonRpcMessage, ServerNotification, + CancelledNotification, CancelledNotificationMethod, CancelledNotificationParam, + ClientJsonRpcMessage, ClientNotification, ClientRequest, JsonRpcNotification, + JsonRpcRequest, Notification, ProgressNotificationParam, ProgressToken, RequestId, + ServerJsonRpcMessage, ServerNotification, }, transport::{ WorkerTransport, @@ -92,6 +95,13 @@ impl SessionManager for LocalSessionManager { let handle = sessions .get(id) .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?; + // Remember the request id so the in-flight request can be cancelled if + // the client disconnects before its response is delivered (issue #857). + // Only requests reach `create_stream`; other message kinds carry no id. + let request_id = match &message { + ClientJsonRpcMessage::Request(request) => Some(request.id.clone()), + _ => None, + }; let receiver = handle.establish_request_wise_channel().await?; let http_request_id = receiver.http_request_id; handle.push_message(message, http_request_id).await?; @@ -103,7 +113,8 @@ impl SessionManager for LocalSessionManager { }; ServerSseMessage::priming(event_id, retry) }); - Ok(futures::stream::iter(priming).chain(ReceiverStream::new(receiver.inner))) + let stream = futures::stream::iter(priming).chain(ReceiverStream::new(receiver.inner)); + Ok(CancelOnDisconnect::new(stream, handle.clone(), request_id)) } async fn create_standalone_stream( @@ -159,6 +170,91 @@ impl SessionManager for LocalSessionManager { } } +/// SSE response body wrapper that cancels the in-flight request when the client +/// disconnects before the response is delivered. +/// +/// The streamable-HTTP server returns a request-wise SSE stream as the HTTP +/// response body for each request. If the client drops the connection while a +/// handler is still running (Ctrl-C, network drop, read timeout), the runtime +/// drops this stream. Nothing else feeds that drop back into the per-request +/// cancellation token (`RequestContext::ct`), so a long-running handler would +/// otherwise run to completion. This guard detects the drop-before-completion +/// and injects a synthetic `notifications/cancelled` for the request, reusing +/// the same path as an explicit client cancellation. +/// +/// See . +struct CancelOnDisconnect { + inner: S, + handle: LocalSessionHandle, + /// The request id to cancel on disconnect, taken once a cancellation is + /// issued. `None` when there is nothing to cancel (non-request messages). + request_id: Option, + /// Set once the inner stream terminates (response delivered, or the stream + /// was closed for a reconnection handoff), so drop does not cancel a request + /// that already completed. + completed: bool, +} + +impl CancelOnDisconnect { + fn new(inner: S, handle: LocalSessionHandle, request_id: Option) -> Self { + Self { + inner, + handle, + request_id, + completed: false, + } + } +} + +impl + Unpin> Stream for CancelOnDisconnect { + type Item = ServerSseMessage; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let poll = Pin::new(&mut this.inner).poll_next(cx); + if matches!(poll, Poll::Ready(None)) { + // Stream ended normally: the response was delivered (the worker + // closes the request-wise sender on completion) or the stream was + // handed off for reconnection. Either way, do not cancel. + this.completed = true; + } + poll + } +} + +impl Drop for CancelOnDisconnect { + fn drop(&mut self) { + if self.completed { + return; + } + let Some(request_id) = self.request_id.take() else { + return; + }; + // The response stream was dropped before completing, i.e. the client + // disconnected. Inject a synthetic cancellation so the service fires the + // request's `RequestContext::ct`, mirroring an explicit + // `notifications/cancelled` from the client. + let cancelled = CancelledNotification { + params: CancelledNotificationParam::new( + Some(request_id), + Some("client disconnected".to_string()), + ), + method: CancelledNotificationMethod, + extensions: Default::default(), + }; + let message = + ClientJsonRpcMessage::notification(ClientNotification::CancelledNotification(cancelled)); + // `Drop` cannot await; deliver on the session's control channel without + // blocking. If the channel is momentarily full the request is not + // cancelled early, but the session keep-alive/idle timeout still reaps + // it — no worse than before this fix. + let _ = self.handle.event_tx.try_send(SessionEvent::ClientMessage { + message, + http_request_id: None, + }); + } +} + /// `/request_id>` #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct EventId { diff --git a/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs b/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs new file mode 100644 index 00000000..7931940f --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs @@ -0,0 +1,214 @@ +#![cfg(not(feature = "local"))] +//! Regression test for issue #857: when a streamable-HTTP client disconnects +//! (TCP close) while a tool handler is still running, the in-flight request +//! must be cancelled so the handler's `RequestContext::ct` fires, instead of +//! letting the handler run to completion. + +use std::{sync::Arc, time::Duration}; + +use rmcp::{ + ErrorData as McpError, RoleServer, ServerHandler, + model::{CallToolRequestParams, CallToolResult, ContentBlock, ServerCapabilities, ServerInfo}, + service::RequestContext, + transport::streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, + session::{SessionId, local::LocalSessionManager}, + }, +}; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; + +/// A minimal server exposing: +/// - `long_sleep`: runs until its per-request cancellation token fires (then +/// notifies `cancelled`) or 30s elapses. +/// - `quick`: returns immediately. +#[derive(Clone)] +struct DisconnectServer { + cancelled: Arc, +} + +impl ServerHandler for DisconnectServer { + fn get_info(&self) -> ServerInfo { + ServerInfo { + capabilities: ServerCapabilities::builder().enable_tools().build(), + ..Default::default() + } + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + match request.name.as_ref() { + "long_sleep" => { + tokio::select! { + _ = context.ct.cancelled() => { + self.cancelled.notify_one(); + Ok(CallToolResult::success(vec![ContentBlock::text("cancelled")])) + } + _ = tokio::time::sleep(Duration::from_secs(30)) => { + Ok(CallToolResult::success(vec![ContentBlock::text( + "ran_to_completion", + )])) + } + } + } + "quick" => Ok(CallToolResult::success(vec![ContentBlock::text( + "quick_done", + )])), + other => Err(McpError::invalid_params( + format!("unknown tool: {other}"), + None, + )), + } + } +} + +async fn spawn_server( + server: DisconnectServer, +) -> (String, CancellationToken, tokio::task::JoinHandle<()>) { + let ct = CancellationToken::new(); + let service: StreamableHttpService = + StreamableHttpService::new( + move || Ok(server.clone()), + Default::default(), + StreamableHttpServerConfig::default() + // Short keep-alive so the server observes the client disconnect + // (a failed keep-alive write) quickly. + .with_sse_keep_alive(Some(Duration::from_millis(100))) + .with_cancellation_token(ct.child_token()), + ); + + let router = axum::Router::new().nest_service("/mcp", service); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let handle = tokio::spawn({ + let ct = ct.clone(); + async move { + let _ = axum::serve(listener, router) + .with_graceful_shutdown(async move { ct.cancelled_owned().await }) + .await; + } + }); + (format!("http://{addr}/mcp"), ct, handle) +} + +async fn init_session(client: &reqwest::Client, url: &str) -> SessionId { + let response = client + .post(url) + .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 + .unwrap(); + assert_eq!(response.status(), 200); + let session_id: SessionId = response.headers()["mcp-session-id"] + .to_str() + .unwrap() + .into(); + + let status = client + .post(url) + .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 + .unwrap() + .status(); + assert_eq!(status, 202); + session_id +} + +#[tokio::test] +async fn client_disconnect_cancels_in_flight_request() -> anyhow::Result<()> { + let cancelled = Arc::new(Notify::new()); + let (url, ct, handle) = spawn_server(DisconnectServer { + cancelled: cancelled.clone(), + }) + .await; + + // Disable connection pooling so dropping the response actually closes the + // TCP connection (simulating a client disconnect). + let client = reqwest::Client::builder() + .pool_max_idle_per_host(0) + .build()?; + let session_id = init_session(&client, &url).await; + + let response = client + .post(&url) + .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":"long_sleep","arguments":{}}}"#, + ) + .send() + .await?; + assert_eq!(response.status(), 200); + + // Let the handler start awaiting its cancellation token, then disconnect by + // dropping the streaming response without reading it to completion. + tokio::time::sleep(Duration::from_millis(200)).await; + drop(response); + + // The fix must fire the handler's cancellation token; `long_sleep` then + // returns via its cancel branch and notifies us. Without the fix the + // handler sleeps for 30s and this times out. + tokio::time::timeout(Duration::from_secs(10), cancelled.notified()) + .await + .expect("handler cancellation token should fire after client disconnect"); + + ct.cancel(); + let _ = handle.await; + Ok(()) +} + +#[tokio::test] +async fn normal_tool_response_is_delivered() -> anyhow::Result<()> { + let cancelled = Arc::new(Notify::new()); + let (url, ct, handle) = spawn_server(DisconnectServer { + cancelled: cancelled.clone(), + }) + .await; + + let client = reqwest::Client::new(); + let session_id = init_session(&client, &url).await; + + // A normal, fully-read tool call must still work end-to-end: the disconnect + // guard must forward the response unchanged and not cancel a request that + // completes normally. + let body = client + .post(&url) + .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":"quick","arguments":{}}}"#, + ) + .send() + .await? + .text() + .await?; + + assert!( + body.contains("quick_done"), + "expected tool result in response, got: {body}" + ); + assert!( + body.contains(r#""id":2"#), + "expected response id 2, got: {body}" + ); + + ct.cancel(); + let _ = handle.await; + Ok(()) +} From b6c93facaf68859cec9e08661f8b38845c460811 Mon Sep 17 00:00:00 2001 From: Amey Pawar <138877912+ameyypawar@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:58:32 +0530 Subject: [PATCH 2/3] test(transport): register streamable-http disconnect-cancel test (#857) Add the `[[test]]` entry so the new integration test is only built when its required features (server, transport-streamable-http-server, reqwest) are enabled, matching the other streamable-http test targets. --- crates/rmcp/Cargo.toml | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 9704bfc7..2c4e0bea 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -387,3 +387,12 @@ required-features = [ "transport-streamable-http-client-reqwest", ] path = "tests/test_streamable_http_connection_reuse.rs" + +[[test]] +name = "test_streamable_http_disconnect_cancel" +required-features = [ + "server", + "transport-streamable-http-server", + "reqwest", +] +path = "tests/test_streamable_http_disconnect_cancel.rs" From 10c947ec03841304dd39b9f0842f78418cdd3441 Mon Sep 17 00:00:00 2001 From: Amey Pawar <138877912+ameyypawar@users.noreply.github.com> Date: Fri, 10 Jul 2026 00:35:49 +0530 Subject: [PATCH 3/3] fix(transport): satisfy nightly rustfmt and build test ServerInfo via ::new (#857) - Wrap the synthetic-cancellation construction to match `cargo +nightly fmt`. - `ServerInfo` is `#[non_exhaustive]`; construct it in the test with `ServerInfo::new(..)` instead of a struct literal. Verified locally: `cargo +nightly fmt --all -- --check` clean and the regression test passes (disconnect cancels the in-flight request). --- .../src/transport/streamable_http_server/session/local.rs | 5 +++-- crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs | 5 +---- 2 files changed, 4 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 300f36cf..da89e740 100644 --- a/crates/rmcp/src/transport/streamable_http_server/session/local.rs +++ b/crates/rmcp/src/transport/streamable_http_server/session/local.rs @@ -242,8 +242,9 @@ impl Drop for CancelOnDisconnect { method: CancelledNotificationMethod, extensions: Default::default(), }; - let message = - ClientJsonRpcMessage::notification(ClientNotification::CancelledNotification(cancelled)); + let message = ClientJsonRpcMessage::notification( + ClientNotification::CancelledNotification(cancelled), + ); // `Drop` cannot await; deliver on the session's control channel without // blocking. If the channel is momentarily full the request is not // cancelled early, but the session keep-alive/idle timeout still reaps diff --git a/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs b/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs index 7931940f..36deb5b7 100644 --- a/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs +++ b/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs @@ -29,10 +29,7 @@ struct DisconnectServer { impl ServerHandler for DisconnectServer { fn get_info(&self) -> ServerInfo { - ServerInfo { - capabilities: ServerCapabilities::builder().enable_tools().build(), - ..Default::default() - } + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) } async fn call_tool(