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

A stateless streamable-HTTP request is one-shot (no session, no resumption),
so if the client drops the response before the handler finishes, the request
is terminal and should be cancelled. Previously the handler kept running with
its `RequestContext::ct` never firing, so long-running or destructive tools
could not observe a client disconnect.

Give each stateless request its own cancellation token via
`serve_directly_with_ct` and cancel it when the client disconnects. This covers
both stateless paths: `serve_negotiated_request_directly` (per-request version
negotiation) and the non-negotiated path. In each:

- A disconnect while the handler is still producing its first message cancels
  it. The guard is disarmed once the handler emits anything, so a normal
  response is never cancelled.
- The SSE response stream is wrapped in a guard that cancels the handler if the
  stream is dropped before it ends naturally.
- When a negotiated request replies directly (JSON mode, or a non-OK status),
  the receiver is dropped, so a still-running handler is cancelled rather than
  left emitting into a closed channel. Without this its terminal send fails
  before adding the termination permit and the serve loop parks forever.

Stateful (resumable) mode is intentionally left unchanged: there a disconnect
may be recovered via `Last-Event-ID`, so cancelling on disconnect would break
resumption.

Adds a regression test covering both stateless sub-modes (SSE and JSON).

Author: ameyypawar

crates/rmcp/Cargo.toml

@@ -426,3 +426,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;
 
@@ -22,7 +30,7 @@ use crate::{
         ProtocolVersion, RequestId, ServerJsonRpcMessage,
     },
     serve_server,
-    service::serve_directly,
+    service::serve_directly_with_ct,
     transport::{
         OneshotTransport, TransportAdapterIdentity,
         common::{
@@ -852,14 +860,28 @@ where
         request.request.extensions_mut().insert(parts);
         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), as in the
+        // non-negotiated stateless path below.
+        let request_ct = CancellationToken::new();
+        let service = serve_directly_with_ct(service, transport, peer_info, request_ct.clone());
         tokio::spawn(async move {
             let _ = service.waiting().await;
         });
 
         let cancel = self.config.cancellation_token.child_token();
+        // Cancel the handler if the client disconnects while it is still
+        // producing its first message (this future is dropped before
+        // `receiver.recv()` completes). Disarmed once the handler emits
+        // anything, so a normal response is never cancelled.
+        let mut disconnect_guard = Some(request_ct.clone().drop_guard());
         let first = tokio::select! {
-            message = receiver.recv() => message,
+            message = receiver.recv() => {
+                if let Some(guard) = disconnect_guard.take() {
+                    guard.disarm();
+                }
+                message
+            }
             _ = cancel.cancelled() => None,
         }
         .ok_or_else(|| {
@@ -870,17 +892,26 @@ where
         })?;
 
         if self.config.json_response || jsonrpc_http_status(&first) != http::StatusCode::OK {
+            // This message is the whole reply, so `receiver` is dropped here and
+            // anything the handler emits afterwards is undeliverable. Cancel it so
+            // a still-running handler stops instead of running on unobserved: its
+            // terminal `send` would otherwise fail before adding the termination
+            // permit, leaving the serve loop parked forever. A no-op when the
+            // handler already completed.
+            request_ct.cancel();
             return jsonrpc_message_response(first, true);
         }
 
+        // The handler may still be streaming, so guard the response: dropping it
+        // (client disconnect) must cancel the handler.
         let stream = futures::stream::once(async move { first })
             .chain(ReceiverStream::new(receiver))
             .map(|message| {
                 tracing::trace!(?message);
                 ServerSseMessage::from_message(message)
             });
         Ok(sse_stream_response(
-            stream,
+            CancelOnDisconnect::new(stream, request_ct),
             self.config.sse_keep_alive,
             self.config.cancellation_token.child_token(),
         ))
@@ -1544,7 +1575,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;
@@ -1554,8 +1591,19 @@ where
                         // emits an intermediate notification or request, preserve
                         // the complete message sequence by falling back to SSE.
                         let cancel = self.config.cancellation_token.child_token();
+                        // Cancel the handler if the client disconnects while it is
+                        // still producing its first message (this future is dropped
+                        // before `receiver.recv()` completes). Disarmed once the
+                        // handler emits anything, so a normal response is never
+                        // cancelled.
+                        let mut disconnect_guard = Some(request_ct.clone().drop_guard());
                         let Some(message) = (tokio::select! {
-                            res = receiver.recv() => res,
+                            res = receiver.recv() => {
+                                if let Some(guard) = disconnect_guard.take() {
+                                    guard.disarm();
+                                }
+                                res
+                            }
                             _ = cancel.cancelled() => None,
                         }) else {
                             return Err(internal_error_response("empty response")(
@@ -1579,6 +1627,9 @@ where
                                 .body(Full::new(Bytes::from(body)).boxed())
                                 .expect("valid response"))
                         } else {
+                            // The handler emitted an intermediate message and is still
+                            // running, so guard the streamed sequence too: dropping it
+                            // (client disconnect) must cancel the handler.
                             let first = futures::stream::once(async move {
                                 ServerSseMessage::from_message(message)
                             });
@@ -1587,17 +1638,19 @@ where
                                 ServerSseMessage::from_message(message)
                             });
                             Ok(sse_stream_response(
-                                first.chain(remaining),
+                                CancelOnDisconnect::new(first.chain(remaining), request_ct),
                                 self.config.sse_keep_alive,
                                 self.config.cancellation_token.child_token(),
                             ))
                         }
                     } 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,
@@ -1680,3 +1733,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
+    }
+}