From 9092067452978f693faecb7a27f02ce8aef053eb Mon Sep 17 00:00:00 2001 From: Amey Pawar Date: Mon, 13 Jul 2026 00:05:06 +0530 Subject: [PATCH] fix(transport): cancel in-flight request on stateless streamable-HTTP client disconnect (#857) 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). --- crates/rmcp/Cargo.toml | 9 + .../transport/streamable_http_server/tower.rs | 122 ++++++++++- .../test_streamable_http_disconnect_cancel.rs | 196 ++++++++++++++++++ 3 files changed, 317 insertions(+), 10 deletions(-) create mode 100644 crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 8989eb06..280cfd37 100644 --- a/crates/rmcp/Cargo.toml +++ b/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" diff --git a/crates/rmcp/src/transport/streamable_http_server/tower.rs b/crates/rmcp/src/transport/streamable_http_server/tower.rs index 1994d554..117912e4 100644 --- a/crates/rmcp/src/transport/streamable_http_server/tower.rs +++ b/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::::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,9 +892,18 @@ 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| { @@ -880,7 +911,7 @@ where 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::::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 { + #[pin] + inner: S, + ct: Option, + } + impl PinnedDrop for CancelOnDisconnect { + fn drop(this: Pin<&mut Self>) { + let this = this.project(); + if let Some(ct) = this.ct.take() { + ct.cancel(); + } + } + } +} + +impl CancelOnDisconnect { + fn new(inner: S, ct: CancellationToken) -> Self { + Self { + inner, + ct: Some(ct), + } + } +} + +impl Stream for CancelOnDisconnect { + type Item = S::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + 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 + } +} 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..73d056fa --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs @@ -0,0 +1,196 @@ +#![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, CallToolResponse, 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, + cancelled: Arc, +} + +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, + ) -> Result { + 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")]).into()) + } + _ = tokio::time::sleep(Duration::from_secs(30)) => { + Ok(CallToolResult::success(vec![ContentBlock::text("ran_to_completion")]).into()) + } + } + } +} + +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, + cancelled: Arc, +} + +async fn spawn_stateless_server(json_response: bool) -> anyhow::Result { + 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 = + 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(()) +}