modelcontextprotocol/rust-sdk
Watch
1
0
Fork
You've already forked rust-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

fix(transport-sse-server): cleanup on connection drop (#165)

Browse source
This commit is contained in:
Aaron 2025-05-12 02:15:34 -04:00 • committed by GitHub
commit 030b6f0858
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 15 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

15
crates/rmcp/src/transport/sse_server.rs
View file

@ -91,6 +91,8 @@ async fn sse_handler(
use tokio_util::sync::PollSender; use tokio_util::sync::PollSender;
let (from_client_tx, from_client_rx) = tokio::sync::mpsc::channel(64); 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, to_client_rx) = tokio::sync::mpsc::channel(64);
let to_client_tx_clone = to_client_tx.clone();
app.txs app.txs
.write() .write()
.await .await
@ -125,6 +127,19 @@ async fn sse_handler(
Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e)), 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))) Ok(Sse::new(stream).keep_alive(KeepAlive::new().interval(ping_interval)))
} }