Skip to content

Commit 030b6f0

Browse files
authored
fix(transport-sse-server): cleanup on connection drop (#165)
1 parent 6de444c commit 030b6f0

File tree

1 file changed

+15
-0
lines changed

1 file changed

+15
-0
lines changed

crates/rmcp/src/transport/sse_server.rs

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,8 @@ async fn sse_handler(
9191
use tokio_util::sync::PollSender;
9292
let (from_client_tx, from_client_rx) = tokio::sync::mpsc::channel(64);
9393
let (to_client_tx, to_client_rx) = tokio::sync::mpsc::channel(64);
94+
let to_client_tx_clone = to_client_tx.clone();
95+
9496
app.txs
9597
.write()
9698
.await
@@ -125,6 +127,19 @@ async fn sse_handler(
125127
Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e)),
126128
}
127129
}));
130+
131+
tokio::spawn(async move {
132+
// Wait for connection closure
133+
to_client_tx_clone.closed().await;
134+
135+
// Clean up session
136+
let session_id = session.clone();
137+
let tx_store = app.txs.clone();
138+
let mut txs = tx_store.write().await;
139+
txs.remove(&session_id);
140+
tracing::debug!(%session_id, "Closed session and cleaned up resources");
141+
});
142+
128143
Ok(Sse::new(stream).keep_alive(KeepAlive::new().interval(ping_interval)))
129144
}
130145

0 commit comments

Comments
 (0)