From aaec423ce2918dc4b57ea5aa761e9bd21acbb07d 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: - SSE mode: a guard on the response stream fires the token if the stream is dropped before it ends naturally. - JSON mode: a drop-guard fires the token if the request future is dropped before the handler produces its response. 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. --- crates/rmcp/Cargo.toml | 9 + .../transport/streamable_http_server/tower.rs | 87 +++++++- .../test_streamable_http_disconnect_cancel.rs | 193 ++++++++++++++++++ 3 files changed, 283 insertions(+), 6 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 e3e9ff0d..d2ca3fe3 100644 --- a/crates/rmcp/Cargo.toml +++ b/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" diff --git a/crates/rmcp/src/transport/streamable_http_server/tower.rs b/crates/rmcp/src/transport/streamable_http_server/tower.rs index 22be7379..9307675f 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; @@ -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::::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 { + #[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..b8accd9b --- /dev/null +++ b/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, + 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")])) + } + _ = 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, + 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(()) +}