From 6fab9f6b23243a1f2782b40be046f7f50636f2f2 Mon Sep 17 00:00:00 2001 From: Ali Hashemi Date: Sat, 16 Aug 2025 17:03:15 -0300 Subject: [PATCH] fix: abort keep-alive task when transport is removed --- crates/rust-mcp-sdk/src/mcp_runtimes/server_runtime.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/crates/rust-mcp-sdk/src/mcp_runtimes/server_runtime.rs b/crates/rust-mcp-sdk/src/mcp_runtimes/server_runtime.rs index d1a8a26..69e8b88 100644 --- a/crates/rust-mcp-sdk/src/mcp_runtimes/server_runtime.rs +++ b/crates/rust-mcp-sdk/src/mcp_runtimes/server_runtime.rs @@ -392,7 +392,10 @@ impl ServerRuntime { let transport = self.transport_by_stream(stream_id).await?; let (disconnect_tx, mut disconnect_rx) = oneshot::channel::<()>(); - let _ = transport.keep_alive(ping_interval, disconnect_tx).await; + let abort_alive_task = transport + .keep_alive(ping_interval, disconnect_tx) + .await? + .abort_handle(); // in case there is a payload, we consume it by transport to get processed if let Some(payload) = payload { @@ -429,11 +432,13 @@ impl ServerRuntime { } // close the stream after all messages are sent, unless it is a standalone stream if !stream_id.eq(DEFAULT_STREAM_ID){ + abort_alive_task.abort(); return Ok(()); } } _ = &mut disconnect_rx => { self.remove_transport(stream_id).await?; + abort_alive_task.abort(); // Disconnection detected by keep-alive task return Err(SdkError::connection_closed().into());