File tree Expand file tree Collapse file tree 1 file changed +15
-0
lines changed
crates/rmcp/src/transport Expand file tree Collapse file tree 1 file changed +15
-0
lines changed Original file line number Diff line number Diff line change @@ -91,6 +91,8 @@ async fn sse_handler(
91
91
use tokio_util:: sync:: PollSender ;
92
92
let ( from_client_tx, from_client_rx) = tokio:: sync:: mpsc:: channel ( 64 ) ;
93
93
let ( to_client_tx, to_client_rx) = tokio:: sync:: mpsc:: channel ( 64 ) ;
94
+ let to_client_tx_clone = to_client_tx. clone ( ) ;
95
+
94
96
app. txs
95
97
. write ( )
96
98
. await
@@ -125,6 +127,19 @@ async fn sse_handler(
125
127
Err ( e) => Err ( io:: Error :: new ( io:: ErrorKind :: InvalidData , e) ) ,
126
128
}
127
129
} ) ) ;
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
+
128
143
Ok ( Sse :: new ( stream) . keep_alive ( KeepAlive :: new ( ) . interval ( ping_interval) ) )
129
144
}
130
145
You can’t perform that action at this time.
0 commit comments