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

ameyypawar · ameyypawar · commit 14fc2f39102e · 2026-07-09T15:56:00.000+05:30
When a streamable-HTTP client closes its TCP connection while a tool
handler is still running, the request-wise SSE response stream is
dropped, but nothing fired the per-request cancellation token. Only a
natural response completion or an explicit notifications/cancelled
cancelled it, so a long-running handler ran to completion after the
client had already disconnected.

Wrap the request-wise response stream in a drop-guard. When the stream
is dropped before completing (client disconnect), it injects a synthetic
notifications/cancelled for the request, reusing the existing
cancellation path so the service fires the handler's RequestContext::ct.
Streams that end normally (response delivered, or a reconnection handoff)
do not cancel.

Adds a regression test covering both the disconnect-cancels and the
normal-completion paths. No public API change.
diff --git a/crates/rmcp/src/transport/streamable_http_server/session/local.rs b/crates/rmcp/src/transport/streamable_http_server/session/local.rs
@@ -1,6 +1,8 @@
 use std::{
     collections::{HashMap, HashSet, VecDeque},
     num::ParseIntError,
+    pin::Pin,
+    task::{Context, Poll},
     time::{Duration, Instant},
 };
 
@@ -16,9 +18,10 @@ use tracing::instrument;
 use crate::{
     RoleServer,
     model::{
-        CancelledNotificationParam, ClientJsonRpcMessage, ClientNotification, ClientRequest,
-        JsonRpcNotification, JsonRpcRequest, Notification, ProgressNotificationParam,
-        ProgressToken, RequestId, ServerJsonRpcMessage, ServerNotification,
+        CancelledNotification, CancelledNotificationMethod, CancelledNotificationParam,
+        ClientJsonRpcMessage, ClientNotification, ClientRequest, JsonRpcNotification,
+        JsonRpcRequest, Notification, ProgressNotificationParam, ProgressToken, RequestId,
+        ServerJsonRpcMessage, ServerNotification,
     },
     transport::{
         WorkerTransport,
@@ -92,6 +95,13 @@ impl SessionManager for LocalSessionManager {
         let handle = sessions
             .get(id)
             .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
+        // Remember the request id so the in-flight request can be cancelled if
+        // the client disconnects before its response is delivered (issue #857).
+        // Only requests reach `create_stream`; other message kinds carry no id.
+        let request_id = match &message {
+            ClientJsonRpcMessage::Request(request) => Some(request.id.clone()),
+            _ => None,
+        };
         let receiver = handle.establish_request_wise_channel().await?;
         let http_request_id = receiver.http_request_id;
         handle.push_message(message, http_request_id).await?;
@@ -103,7 +113,8 @@ impl SessionManager for LocalSessionManager {
             };
             ServerSseMessage::priming(event_id, retry)
         });
-        Ok(futures::stream::iter(priming).chain(ReceiverStream::new(receiver.inner)))
+        let stream = futures::stream::iter(priming).chain(ReceiverStream::new(receiver.inner));
+        Ok(CancelOnDisconnect::new(stream, handle.clone(), request_id))
     }
 
     async fn create_standalone_stream(
@@ -159,6 +170,91 @@ impl SessionManager for LocalSessionManager {
     }
 }
 
+/// SSE response body wrapper that cancels the in-flight request when the client
+/// disconnects before the response is delivered.
+///
+/// The streamable-HTTP server returns a request-wise SSE stream as the HTTP
+/// response body for each request. If the client drops the connection while a
+/// handler is still running (Ctrl-C, network drop, read timeout), the runtime
+/// drops this stream. Nothing else feeds that drop back into the per-request
+/// cancellation token (`RequestContext::ct`), so a long-running handler would
+/// otherwise run to completion. This guard detects the drop-before-completion
+/// and injects a synthetic `notifications/cancelled` for the request, reusing
+/// the same path as an explicit client cancellation.
+///
+/// See <https://github.com/modelcontextprotocol/rust-sdk/issues/857>.
+struct CancelOnDisconnect<S> {
+    inner: S,
+    handle: LocalSessionHandle,
+    /// The request id to cancel on disconnect, taken once a cancellation is
+    /// issued. `None` when there is nothing to cancel (non-request messages).
+    request_id: Option<RequestId>,
+    /// Set once the inner stream terminates (response delivered, or the stream
+    /// was closed for a reconnection handoff), so drop does not cancel a request
+    /// that already completed.
+    completed: bool,
+}
+
+impl<S> CancelOnDisconnect<S> {
+    fn new(inner: S, handle: LocalSessionHandle, request_id: Option<RequestId>) -> Self {
+        Self {
+            inner,
+            handle,
+            request_id,
+            completed: false,
+        }
+    }
+}
+
+impl<S: Stream<Item = ServerSseMessage> + Unpin> Stream for CancelOnDisconnect<S> {
+    type Item = ServerSseMessage;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+        let poll = Pin::new(&mut this.inner).poll_next(cx);
+        if matches!(poll, Poll::Ready(None)) {
+            // Stream ended normally: the response was delivered (the worker
+            // closes the request-wise sender on completion) or the stream was
+            // handed off for reconnection. Either way, do not cancel.
+            this.completed = true;
+        }
+        poll
+    }
+}
+
+impl<S> Drop for CancelOnDisconnect<S> {
+    fn drop(&mut self) {
+        if self.completed {
+            return;
+        }
+        let Some(request_id) = self.request_id.take() else {
+            return;
+        };
+        // The response stream was dropped before completing, i.e. the client
+        // disconnected. Inject a synthetic cancellation so the service fires the
+        // request's `RequestContext::ct`, mirroring an explicit
+        // `notifications/cancelled` from the client.
+        let cancelled = CancelledNotification {
+            params: CancelledNotificationParam::new(
+                Some(request_id),
+                Some("client disconnected".to_string()),
+            ),
+            method: CancelledNotificationMethod,
+            extensions: Default::default(),
+        };
+        let message =
+            ClientJsonRpcMessage::notification(ClientNotification::CancelledNotification(cancelled));
+        // `Drop` cannot await; deliver on the session's control channel without
+        // blocking. If the channel is momentarily full the request is not
+        // cancelled early, but the session keep-alive/idle timeout still reaps
+        // it — no worse than before this fix.
+        let _ = self.handle.event_tx.try_send(SessionEvent::ClientMessage {
+            message,
+            http_request_id: None,
+        });
+    }
+}
+
 /// `<index>/request_id>`
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct EventId {
diff --git a/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs b/crates/rmcp/tests/test_streamable_http_disconnect_cancel.rs
@@ -0,0 +1,214 @@
+#![cfg(not(feature = "local"))]
+//! Regression test for issue #857: when a streamable-HTTP client disconnects
+//! (TCP close) while a tool handler is still running, the in-flight request
+//! must be cancelled so the handler's `RequestContext::ct` fires, instead of
+//! letting the handler run to completion.
+
+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::{SessionId, local::LocalSessionManager},
+    },
+};
+use tokio::sync::Notify;
+use tokio_util::sync::CancellationToken;
+
+/// A minimal server exposing:
+/// - `long_sleep`: runs until its per-request cancellation token fires (then
+///   notifies `cancelled`) or 30s elapses.
+/// - `quick`: returns immediately.
+#[derive(Clone)]
+struct DisconnectServer {
+    cancelled: Arc<Notify>,
+}
+
+impl ServerHandler for DisconnectServer {
+    fn get_info(&self) -> ServerInfo {
+        ServerInfo {
+            capabilities: ServerCapabilities::builder().enable_tools().build(),
+            ..Default::default()
+        }
+    }
+
+    async fn call_tool(
+        &self,
+        request: CallToolRequestParams,
+        context: RequestContext<RoleServer>,
+    ) -> Result<CallToolResult, McpError> {
+        match request.name.as_ref() {
+            "long_sleep" => {
+                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",
+                        )]))
+                    }
+                }
+            }
+            "quick" => Ok(CallToolResult::success(vec![ContentBlock::text(
+                "quick_done",
+            )])),
+            other => Err(McpError::invalid_params(
+                format!("unknown tool: {other}"),
+                None,
+            )),
+        }
+    }
+}
+
+async fn spawn_server(
+    server: DisconnectServer,
+) -> (String, CancellationToken, tokio::task::JoinHandle<()>) {
+    let ct = CancellationToken::new();
+    let service: StreamableHttpService<DisconnectServer, LocalSessionManager> =
+        StreamableHttpService::new(
+            move || Ok(server.clone()),
+            Default::default(),
+            StreamableHttpServerConfig::default()
+                // Short keep-alive so the server observes the client disconnect
+                // (a failed keep-alive write) quickly.
+                .with_sse_keep_alive(Some(Duration::from_millis(100)))
+                .with_cancellation_token(ct.child_token()),
+        );
+
+    let router = axum::Router::new().nest_service("/mcp", service);
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let handle = tokio::spawn({
+        let ct = ct.clone();
+        async move {
+            let _ = axum::serve(listener, router)
+                .with_graceful_shutdown(async move { ct.cancelled_owned().await })
+                .await;
+        }
+    });
+    (format!("http://{addr}/mcp"), ct, handle)
+}
+
+async fn init_session(client: &reqwest::Client, url: &str) -> SessionId {
+    let response = client
+        .post(url)
+        .header("Content-Type", "application/json")
+        .header("Accept", "application/json, text/event-stream")
+        .body(
+            r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#,
+        )
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(response.status(), 200);
+    let session_id: SessionId = response.headers()["mcp-session-id"]
+        .to_str()
+        .unwrap()
+        .into();
+
+    let status = client
+        .post(url)
+        .header("Content-Type", "application/json")
+        .header("Accept", "application/json, text/event-stream")
+        .header("mcp-session-id", session_id.to_string())
+        .header("Mcp-Protocol-Version", "2025-06-18")
+        .body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
+        .send()
+        .await
+        .unwrap()
+        .status();
+    assert_eq!(status, 202);
+    session_id
+}
+
+#[tokio::test]
+async fn client_disconnect_cancels_in_flight_request() -> anyhow::Result<()> {
+    let cancelled = Arc::new(Notify::new());
+    let (url, ct, handle) = spawn_server(DisconnectServer {
+        cancelled: cancelled.clone(),
+    })
+    .await;
+
+    // Disable connection pooling so dropping the response actually closes the
+    // TCP connection (simulating a client disconnect).
+    let client = reqwest::Client::builder()
+        .pool_max_idle_per_host(0)
+        .build()?;
+    let session_id = init_session(&client, &url).await;
+
+    let response = client
+        .post(&url)
+        .header("Content-Type", "application/json")
+        .header("Accept", "application/json, text/event-stream")
+        .header("mcp-session-id", session_id.to_string())
+        .header("Mcp-Protocol-Version", "2025-06-18")
+        .body(
+            r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"long_sleep","arguments":{}}}"#,
+        )
+        .send()
+        .await?;
+    assert_eq!(response.status(), 200);
+
+    // Let the handler start awaiting its cancellation token, then disconnect by
+    // dropping the streaming response without reading it to completion.
+    tokio::time::sleep(Duration::from_millis(200)).await;
+    drop(response);
+
+    // The fix must fire the handler's cancellation token; `long_sleep` then
+    // returns via its cancel branch and notifies us. Without the fix the
+    // handler sleeps for 30s and this times out.
+    tokio::time::timeout(Duration::from_secs(10), cancelled.notified())
+        .await
+        .expect("handler cancellation token should fire after client disconnect");
+
+    ct.cancel();
+    let _ = handle.await;
+    Ok(())
+}
+
+#[tokio::test]
+async fn normal_tool_response_is_delivered() -> anyhow::Result<()> {
+    let cancelled = Arc::new(Notify::new());
+    let (url, ct, handle) = spawn_server(DisconnectServer {
+        cancelled: cancelled.clone(),
+    })
+    .await;
+
+    let client = reqwest::Client::new();
+    let session_id = init_session(&client, &url).await;
+
+    // A normal, fully-read tool call must still work end-to-end: the disconnect
+    // guard must forward the response unchanged and not cancel a request that
+    // completes normally.
+    let body = client
+        .post(&url)
+        .header("Content-Type", "application/json")
+        .header("Accept", "application/json, text/event-stream")
+        .header("mcp-session-id", session_id.to_string())
+        .header("Mcp-Protocol-Version", "2025-06-18")
+        .body(
+            r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"quick","arguments":{}}}"#,
+        )
+        .send()
+        .await?
+        .text()
+        .await?;
+
+    assert!(
+        body.contains("quick_done"),
+        "expected tool result in response, got: {body}"
+    );
+    assert!(
+        body.contains(r#""id":2"#),
+        "expected response id 2, got: {body}"
+    );
+
+    ct.cancel();
+    let _ = handle.await;
+    Ok(())
+}
