From 030b6f085892888dd1548ae5a6a7efc1f0c5b753 Mon Sep 17 00:00:00 2001 From: Aaron Date: Mon, 12 May 2025 02:15:34 -0400 Subject: [PATCH] fix(transport-sse-server): cleanup on connection drop (#165) --- crates/rmcp/src/transport/sse_server.rs | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/crates/rmcp/src/transport/sse_server.rs b/crates/rmcp/src/transport/sse_server.rs index bbeecf1..7dce63c 100644 --- a/crates/rmcp/src/transport/sse_server.rs +++ b/crates/rmcp/src/transport/sse_server.rs @@ -91,6 +91,8 @@ async fn sse_handler( use tokio_util::sync::PollSender; let (from_client_tx, from_client_rx) = tokio::sync::mpsc::channel(64); let (to_client_tx, to_client_rx) = tokio::sync::mpsc::channel(64); + let to_client_tx_clone = to_client_tx.clone(); + app.txs .write() .await @@ -125,6 +127,19 @@ async fn sse_handler( Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e)), } })); + + tokio::spawn(async move { + // Wait for connection closure + to_client_tx_clone.closed().await; + + // Clean up session + let session_id = session.clone(); + let tx_store = app.txs.clone(); + let mut txs = tx_store.write().await; + txs.remove(&session_id); + tracing::debug!(%session_id, "Closed session and cleaned up resources"); + }); + Ok(Sse::new(stream).keep_alive(KeepAlive::new().interval(ping_interval))) }