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(()) +}