fix(transport): cancel in-flight request on streamable-HTTP client disconnect (#857)

crates/rmcp/Cargo.toml

@@ -382,3 +382,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"

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

@@ -1,12 +1,20 @@
 use std::{
-    borrow::Cow, collections::HashMap, convert::Infallible, fmt::Display, sync::Arc, time::Duration,
+    borrow::Cow,
+    collections::HashMap,
+    convert::Infallible,
+    fmt::Display,
+    pin::Pin,
+    sync::Arc,
+    task::{Context, Poll},
+    time::Duration,
 };
 
 use bytes::Bytes;
-use futures::{StreamExt, future::BoxFuture};
+use futures::{Stream, StreamExt, future::BoxFuture};
 use http::{HeaderMap, Method, Request, Response, header::ALLOW};
 use http_body::Body;
 use http_body_util::{BodyExt, Full, combinators::BoxBody};
+use pin_project_lite::pin_project;
 use tokio_stream::wrappers::ReceiverStream;
 use tokio_util::sync::CancellationToken;
 
@@ -21,7 +29,7 @@ use crate::{
         InitializedNotification, JsonRpcError, ProtocolVersion, RequestId,
     },
     serve_server,
-    service::serve_directly,
+    service::serve_directly_with_ct,
     transport::{
         OneshotTransport, TransportAdapterIdentity,
         common::{
@@ -1250,7 +1258,13 @@ where
                     request.request.extensions_mut().insert(part);
                     let (transport, mut receiver) =
                         OneshotTransport::<RoleServer>::new(ClientJsonRpcMessage::Request(request));
-                    let service = serve_directly(service, transport, peer_info);
+                    // Give this stateless request its own cancellation token so a
+                    // client disconnect can cancel the in-flight handler (#857). A
+                    // stateless request is one-shot (no session, no resumption), so a
+                    // dropped response is terminal and safe to cancel.
+                    let request_ct = CancellationToken::new();
+                    let service =
+                        serve_directly_with_ct(service, transport, peer_info, request_ct.clone());
                     tokio::spawn(async move {
                         // on service created
                         let _ = service.waiting().await;
@@ -1260,8 +1274,18 @@ where
                         // application/json, eliminating SSE framing overhead.
                         // Allowed by MCP Streamable HTTP spec (2025-06-18).
                         let cancel = self.config.cancellation_token.child_token();
+                        // Cancel the handler only if the client disconnects while it
+                        // is still producing its first message (this future is dropped
+                        // before `receiver.recv()` completes). Once the handler emits
+                        // anything we disarm, so a normal response is never cancelled.
+                        let mut disconnect_guard = Some(request_ct.drop_guard());
                         match tokio::select! {
-                            res = receiver.recv() => res,
+                            res = receiver.recv() => {
+                                if let Some(guard) = disconnect_guard.take() {
+                                    guard.disarm();
+                                }
+                                res
+                            }
                             _ = cancel.cancelled() => None,
                         } {
                             Some(message) => {
@@ -1283,11 +1307,13 @@ where
                             )),
                         }
                     } else {
-                        // SSE mode (default): original behaviour preserved unchanged
+                        // SSE mode (default): cancel the handler if the client
+                        // disconnects (drops the response stream) before it completes.
                         let stream = ReceiverStream::new(receiver).map(|message| {
                             tracing::trace!(?message);
                             ServerSseMessage::from_message(message)
                         });
+                        let stream = CancelOnDisconnect::new(stream, request_ct);
                         Ok(sse_stream_response(
                             stream,
                             self.config.sse_keep_alive,
@@ -1370,3 +1396,52 @@ where
         })
     }
 }
+
+pin_project! {
+    /// Wraps a stateless SSE response stream so a client disconnect cancels the
+    /// in-flight request.
+    ///
+    /// A stateless streamable-HTTP request is one-shot: it has no session and no
+    /// resumption, so a dropped response stream means the client is gone for
+    /// good. When the stream is dropped *before* it ends naturally, the request's
+    /// cancellation token is fired, which stops the dedicated `serve_directly`
+    /// loop and cancels the handler's `RequestContext::ct` (see #857). If the
+    /// stream ends naturally (the request completed), the guard is disarmed so
+    /// normal completion cancels nothing.
+    struct CancelOnDisconnect<S> {
+        #[pin]
+        inner: S,
+        ct: Option<CancellationToken>,
+    }
+    impl<S> PinnedDrop for CancelOnDisconnect<S> {
+        fn drop(this: Pin<&mut Self>) {
+            let this = this.project();
+            if let Some(ct) = this.ct.take() {
+                ct.cancel();
+            }
+        }
+    }
+}
+
+impl<S> CancelOnDisconnect<S> {
+    fn new(inner: S, ct: CancellationToken) -> Self {
+        Self {
+            inner,
+            ct: Some(ct),
+        }
+    }
+}
+
+impl<S: Stream> Stream for CancelOnDisconnect<S> {
+    type Item = S::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+        let polled = this.inner.poll_next(cx);
+        if let Poll::Ready(None) = &polled {
+            // Ended naturally: the request completed, so don't cancel on drop.
+            *this.ct = None;
+        }
+        polled
+    }
+}

crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs

@@ -0,0 +1,193 @@
+#![cfg(all(
+    feature = "server",
+    feature = "transport-streamable-http-server",
+    feature = "reqwest",
+    not(feature = "local")
+))]
+
+//! Regression test for #857: when a stateless streamable-HTTP client disconnects
+//! (drops the response) while a tool handler is still awaiting, the per-request
+//! `RequestContext::ct` should fire so the handler can cancel cooperatively.
+//!
+//! Stateless requests are one-shot (no session, no resumption), so a dropped
+//! response is terminal and safe to cancel — unlike the stateful/resumable path,
+//! where a disconnect may be recovered via `Last-Event-ID`.
+
+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::local::LocalSessionManager,
+    },
+};
+use tokio::sync::Notify;
+use tokio_util::sync::CancellationToken;
+
+#[derive(Clone)]
+struct CancelProbe {
+    started: Arc<Notify>,
+    cancelled: Arc<Notify>,
+}
+
+impl ServerHandler for CancelProbe {
+    #[allow(deprecated)]
+    fn get_info(&self) -> ServerInfo {
+        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
+    }
+
+    async fn call_tool(
+        &self,
+        _request: CallToolRequestParams,
+        context: RequestContext<RoleServer>,
+    ) -> Result<CallToolResult, McpError> {
+        self.started.notify_one();
+        // Wait until the per-request cancellation token fires, or give up after a
+        // generous timeout so a buggy build fails via the outer assertion rather
+        // than hanging the test.
+        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")]))
+            }
+        }
+    }
+}
+
+const CALL_BODY: &str = r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"wait_for_cancel","arguments":{}}}"#;
+
+struct TestServer {
+    url: String,
+    server_ct: CancellationToken,
+    started: Arc<Notify>,
+    cancelled: Arc<Notify>,
+}
+
+async fn spawn_stateless_server(json_response: bool) -> anyhow::Result<TestServer> {
+    let started = Arc::new(Notify::new());
+    let cancelled = Arc::new(Notify::new());
+    let probe = CancelProbe {
+        started: started.clone(),
+        cancelled: cancelled.clone(),
+    };
+
+    let server_ct = CancellationToken::new();
+    let config = StreamableHttpServerConfig::default()
+        .with_stateful_mode(false)
+        .with_json_response(json_response)
+        // A short keep-alive lets the SSE server notice a dropped connection
+        // quickly (hyper only observes the disconnect on its next write).
+        .with_sse_keep_alive(Some(Duration::from_millis(100)))
+        .with_cancellation_token(server_ct.child_token());
+
+    let service: StreamableHttpService<CancelProbe, LocalSessionManager> =
+        StreamableHttpService::new(
+            move || Ok(probe.clone()),
+            Arc::new(LocalSessionManager::default()),
+            config,
+        );
+    let router = axum::Router::new().nest_service("/mcp", service);
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+    let addr = listener.local_addr()?;
+    tokio::spawn({
+        let ct = server_ct.clone();
+        async move {
+            let _ = axum::serve(listener, router)
+                .with_graceful_shutdown(async move { ct.cancelled_owned().await })
+                .await;
+        }
+    });
+
+    Ok(TestServer {
+        url: format!("http://{addr}/mcp"),
+        server_ct,
+        started,
+        cancelled,
+    })
+}
+
+/// SSE mode: the response is a stream; dropping it (client disconnect) must fire
+/// the handler's cancellation token.
+#[tokio::test]
+async fn stateless_sse_client_disconnect_cancels_request() -> anyhow::Result<()> {
+    let server = spawn_stateless_server(false).await?;
+    let client = reqwest::Client::builder()
+        .pool_max_idle_per_host(0)
+        .build()?;
+
+    // A single self-contained tools/call (no session, no initialize handshake).
+    let call = client
+        .post(&server.url)
+        .header("Content-Type", "application/json")
+        .header("Accept", "application/json, text/event-stream")
+        .header("MCP-Protocol-Version", "2025-03-26")
+        .body(CALL_BODY)
+        .send()
+        .await?;
+    assert!(
+        call.status().is_success(),
+        "tools/call failed: {:?}",
+        call.status()
+    );
+
+    tokio::time::timeout(Duration::from_secs(5), server.started.notified())
+        .await
+        .expect("tool handler should start");
+
+    // Client disconnects mid-call: drop the streaming response (and the client).
+    drop(call);
+    drop(client);
+
+    tokio::time::timeout(Duration::from_secs(10), server.cancelled.notified())
+        .await
+        .expect("RequestContext::ct should fire after client disconnect (SSE)");
+
+    server.server_ct.cancel();
+    Ok(())
+}
+
+/// JSON-direct mode: the server holds the connection open awaiting the single
+/// response. A client that disconnects while the handler is running must still
+/// fire the handler's cancellation token.
+#[tokio::test]
+async fn stateless_json_client_disconnect_cancels_request() -> anyhow::Result<()> {
+    let server = spawn_stateless_server(true).await?;
+    let client = reqwest::Client::builder()
+        .pool_max_idle_per_host(0)
+        .build()?;
+
+    // In JSON mode the server does not respond until the handler completes, so
+    // the request stays pending; drive it from a task we can abort to disconnect.
+    let url = server.url.clone();
+    let req_task = tokio::spawn(async move {
+        let _ = client
+            .post(&url)
+            .header("Content-Type", "application/json")
+            .header("Accept", "application/json, text/event-stream")
+            .header("MCP-Protocol-Version", "2025-03-26")
+            .body(CALL_BODY)
+            .send()
+            .await;
+        // Keep the client alive until the request future is dropped by abort().
+        drop(client);
+    });
+
+    tokio::time::timeout(Duration::from_secs(5), server.started.notified())
+        .await
+        .expect("tool handler should start");
+
+    // Client disconnects: abort the in-flight request, closing the connection.
+    req_task.abort();
+
+    tokio::time::timeout(Duration::from_secs(10), server.cancelled.notified())
+        .await
+        .expect("RequestContext::ct should fire after client disconnect (JSON)");
+
+    server.server_ct.cancel();
+    Ok(())
+}