Skip to content

Commit 45f2f72

Browse files
jstar0DaleSeo
jstar0
and
DaleSeo
authored
fix: fail orphaned streamable HTTP responses on reinit (#914)
* fix: fail orphaned streamable HTTP responses on reinit * fix: update crates/rmcp/src/transport/streamable_http_client.rs --------- Co-authored-by: Dale Seo <5466341+DaleSeo@users.noreply.github.com>
1 parent 95490fa commit 45f2f72

2 files changed

Lines changed: 343 additions & 7 deletions

File tree

‎crates/rmcp/src/transport/streamable_http_client.rs‎

Lines changed: 122 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,9 @@
1-
use std::{borrow::Cow, collections::HashMap, sync::Arc, time::Duration};
1+
use std::{
2+
borrow::Cow,
3+
collections::{HashMap, HashSet},
4+
sync::Arc,
5+
time::Duration,
6+
};
27

38
use futures::{Stream, StreamExt, future::BoxFuture, stream::BoxStream};
49
use http::{HeaderName, HeaderValue};
@@ -12,8 +17,8 @@ use super::common::client_side_sse::{ExponentialBackoff, SseRetryPolicy, SseStre
1217
use crate::{
1318
RoleClient,
1419
model::{
15-
ClientJsonRpcMessage, ClientNotification, InitializedNotification, ServerJsonRpcMessage,
16-
ServerResult,
20+
ClientJsonRpcMessage, ClientNotification, ErrorData, InitializedNotification, RequestId,
21+
ServerJsonRpcMessage, ServerResult,
1722
},
1823
transport::{
1924
common::client_side_sse::SseAutoReconnectStream,
@@ -298,6 +303,79 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
298303
}
299304

300305
impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
306+
fn client_request_id(message: &ClientJsonRpcMessage) -> Option<RequestId> {
307+
match message {
308+
ClientJsonRpcMessage::Request(request) => Some(request.id.clone()),
309+
_ => None,
310+
}
311+
}
312+
313+
fn server_response_id(message: &ServerJsonRpcMessage) -> Option<&RequestId> {
314+
match message {
315+
ServerJsonRpcMessage::Response(response) => Some(&response.id),
316+
ServerJsonRpcMessage::Error(error) => error.id.as_ref(),
317+
_ => None,
318+
}
319+
}
320+
321+
fn mark_stream_response_pending(
322+
pending_stream_response_ids: &mut HashSet<RequestId>,
323+
request_id: Option<RequestId>,
324+
) {
325+
if let Some(request_id) = request_id {
326+
pending_stream_response_ids.insert(request_id);
327+
}
328+
}
329+
330+
fn clear_stream_response_pending(
331+
pending_stream_response_ids: &mut HashSet<RequestId>,
332+
message: &ServerJsonRpcMessage,
333+
) {
334+
if let Some(id) = Self::server_response_id(message) {
335+
pending_stream_response_ids.remove(id);
336+
}
337+
}
338+
339+
async fn drain_queued_stream_messages(
340+
sse_worker_rx: &mut tokio::sync::mpsc::Receiver<ServerJsonRpcMessage>,
341+
context: &mut super::worker::WorkerContext<Self>,
342+
pending_stream_response_ids: &mut HashSet<RequestId>,
343+
) -> Result<(), WorkerQuitReason<StreamableHttpError<C::Error>>> {
344+
loop {
345+
match sse_worker_rx.try_recv() {
346+
Ok(message) => {
347+
Self::clear_stream_response_pending(pending_stream_response_ids, &message);
348+
context.send_to_handler(message).await?;
349+
}
350+
Err(tokio::sync::mpsc::error::TryRecvError::Empty) => return Ok(()),
351+
Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => return Ok(()),
352+
}
353+
}
354+
}
355+
356+
async fn fail_pending_stream_responses(
357+
context: &mut super::worker::WorkerContext<Self>,
358+
pending_stream_response_ids: &mut HashSet<RequestId>,
359+
) -> Result<(), WorkerQuitReason<StreamableHttpError<C::Error>>> {
360+
if pending_stream_response_ids.is_empty() {
361+
return Ok(());
362+
}
363+
364+
let pending_ids = std::mem::take(pending_stream_response_ids);
365+
for id in pending_ids {
366+
context
367+
.send_to_handler(ServerJsonRpcMessage::error(
368+
ErrorData::internal_error(
369+
"streamable HTTP session was re-initialized before the response arrived",
370+
None,
371+
),
372+
Some(id),
373+
))
374+
.await?;
375+
}
376+
Ok(())
377+
}
378+
301379
/// Convert a raw SSE stream into a JSON-RPC message stream without
302380
/// reconnection logic.
303381
fn raw_sse_to_jsonrpc(
@@ -557,6 +635,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
557635
StreamResult(Result<(), StreamableHttpError<E>>),
558636
}
559637
let mut streams = tokio::task::JoinSet::new();
638+
let mut pending_stream_response_ids = HashSet::new();
560639
if let Some(session_id) = &session_id {
561640
let client = self.client.clone();
562641
let uri = config.uri.clone();
@@ -646,6 +725,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
646725
match event {
647726
Event::ClientMessage(send_request) => {
648727
let WorkerSendRequest { message, responder } = send_request;
728+
let request_id = Self::client_request_id(&message);
649729
// Pass a clone to the first attempt so `message` is retained for a
650730
// potential re-init retry. `post_message` takes ownership and the
651731
// trait cannot be changed, so the clone is unavoidable.
@@ -679,9 +759,26 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
679759
.await
680760
{
681761
Ok((new_session_id, new_protocol_headers)) => {
682-
// Old streams hold the stale session ID; abort them
683-
// so the new standalone SSE stream takes over.
762+
// Old streams hold the stale session ID. Stop them first
763+
// so no late stale-session messages can arrive after the
764+
// pending requests below are completed.
684765
streams.abort_all();
766+
while streams.join_next().await.is_some() {}
767+
768+
// Forward any already queued response messages and fail
769+
// the remaining accepted requests so callers do not wait
770+
// forever for responses that can no longer arrive.
771+
Self::drain_queued_stream_messages(
772+
&mut sse_worker_rx,
773+
&mut context,
774+
&mut pending_stream_response_ids,
775+
)
776+
.await?;
777+
Self::fail_pending_stream_responses(
778+
&mut context,
779+
&mut pending_stream_response_ids,
780+
)
781+
.await?;
685782

686783
session_id = new_session_id;
687784
protocol_headers = new_protocol_headers;
@@ -765,6 +862,10 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
765862
match retry_response {
766863
Err(e) => Err(e),
767864
Ok(StreamableHttpPostResponse::Accepted) => {
865+
Self::mark_stream_response_pending(
866+
&mut pending_stream_response_ids,
867+
request_id,
868+
);
768869
tracing::trace!(
769870
"client message accepted after re-init"
770871
);
@@ -775,6 +876,10 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
775876
Ok(())
776877
}
777878
Ok(StreamableHttpPostResponse::Sse(stream, ..)) => {
879+
Self::mark_stream_response_pending(
880+
&mut pending_stream_response_ids,
881+
request_id,
882+
);
778883
streams.spawn(Self::execute_sse_stream(
779884
Self::raw_sse_to_jsonrpc(stream),
780885
sse_worker_tx.clone(),
@@ -792,6 +897,10 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
792897
}
793898
Err(e) => Err(e),
794899
Ok(StreamableHttpPostResponse::Accepted) => {
900+
Self::mark_stream_response_pending(
901+
&mut pending_stream_response_ids,
902+
request_id,
903+
);
795904
tracing::trace!("client message accepted");
796905
Ok(())
797906
}
@@ -800,6 +909,10 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
800909
Ok(())
801910
}
802911
Ok(StreamableHttpPostResponse::Sse(stream, ..)) => {
912+
Self::mark_stream_response_pending(
913+
&mut pending_stream_response_ids,
914+
request_id,
915+
);
803916
streams.spawn(Self::execute_sse_stream(
804917
Self::raw_sse_to_jsonrpc(stream),
805918
sse_worker_tx.clone(),
@@ -813,6 +926,10 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
813926
let _ = responder.send(send_result);
814927
}
815928
Event::ServerMessage(json_rpc_message) => {
929+
Self::clear_stream_response_pending(
930+
&mut pending_stream_response_ids,
931+
&json_rpc_message,
932+
);
816933
// send the message to the handler
817934
if let Err(e) = context.send_to_handler(json_rpc_message).await {
818935
break 'main_loop Err(e);

0 commit comments

Comments
 (0)