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(http): drain SSE stream for connection reuse (#790)

Browse source 
* fix(http): reduce latency on subsequent StreamableHttp calls

* refactor: rely on stream drain for connection reuse

* refactor: clean up comments and naming

* fix: restore pool_max_idle_per_host(0) for Linux
This commit is contained in:
Dale Seo 2026-04-13 16:33:30 -04:00 • committed by GitHub
commit c99903a67a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 210 additions and 69 deletions
Download patch file Download diff file Expand all files Collapse all files

16
crates/rmcp/src/transport/common/reqwest/streamable_http_client.rs
View file

@ -262,7 +262,7 @@ impl StreamableHttpClientTransport<reqwest::Client> {
/// This method requires the `transport-streamable-http-client-reqwest` feature.
pub fn from_uri(uri: impl Into<Arc<str>>) -> Self {
StreamableHttpClientTransport::with_client(
reqwest::Client::default(),
Self::default_http_client(),
StreamableHttpClientTransportConfig {
uri: uri.into(),
auth_header: None,
@ -277,7 +277,19 @@ impl StreamableHttpClientTransport<reqwest::Client> {
///
/// * `config` - The config to use with this transport
pub fn from_config(config: StreamableHttpClientTransportConfig) -> Self {
StreamableHttpClientTransport::with_client(reqwest::Client::default(), config)
StreamableHttpClientTransport::with_client(Self::default_http_client(), config)
}
/// Build the default reqwest client for this transport.
///
/// Disables idle connection pooling to avoid ~40 ms stalls caused by
/// TCP Delayed ACK on Linux when the previous response body was not
/// fully consumed before the pool attempts to reuse the connection.
fn default_http_client() -> reqwest::Client {
reqwest::Client::builder()
.pool_max_idle_per_host(0)
.build()
.expect("failed to build default reqwest client")
}
}

118
crates/rmcp/src/transport/streamable_http_client.rs
View file

@ -298,6 +298,37 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
}
impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
/// Convert a raw SSE stream into a JSON-RPC message stream without
/// reconnection logic.
fn raw_sse_to_jsonrpc(
stream: BoxedSseStream,
) -> impl Stream<Item = Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>> + Send + 'static
{
stream.filter_map(|event| async {
match event {
Err(e) => Some(Err(StreamableHttpError::Sse(e))),
Ok(sse) => {
let is_message =
matches!(sse.event.as_deref(), None | Some("") | Some("message"));
if !is_message {
return None;
}
let data = sse.data?;
if data.trim().is_empty() {
return None;
}
match serde_json::from_str::<ServerJsonRpcMessage>(&data) {
Ok(msg) => Some(Ok(msg)),
Err(e) => {
tracing::debug!("failed to deserialize server message: {e}");
None
}
}
}
}
})
}
async fn execute_sse_stream(
sse_stream: impl Stream<Item = Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>>
+ Send
@ -320,14 +351,23 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
let Some(message) = message.transpose()? else {
break;
};
let is_response = matches!(message, ServerJsonRpcMessage::Response(_));
let is_response = matches!(
message,
ServerJsonRpcMessage::Response(_) | ServerJsonRpcMessage::Error(_)
);
let yield_result = sse_worker_tx.send(message).await;
if yield_result.is_err() {
tracing::trace!("streamable http transport worker dropped, exiting");
break;
}
if close_on_response && is_response {
tracing::debug!("got response, closing sse stream");
tracing::debug!("got response, draining sse stream for connection reuse");
// Consume the remaining stream so the HTTP/1.1 connection
// returns to the pool cleanly.
let _ = tokio::time::timeout(std::time::Duration::from_millis(50), async {
while sse_stream.next().await.is_some() {}
})
.await;
break;
}
}
@ -735,38 +775,12 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
Ok(())
}
Ok(StreamableHttpPostResponse::Sse(stream, ..)) => {
if let Some(sid) = &session_id {
let sse_stream = SseAutoReconnectStream::new(
stream,
StreamableHttpClientReconnect {
client: self.client.clone(),
session_id: sid.clone(),
uri: config.uri.clone(),
auth_header: config.auth_header.clone(),
custom_headers: protocol_headers
.clone(),
},
self.config.retry_config.clone(),
);
streams.spawn(Self::execute_sse_stream(
sse_stream,
sse_worker_tx.clone(),
true,
transport_task_ct.child_token(),
));
} else {
let sse_stream =
SseAutoReconnectStream::never_reconnect(
stream,
StreamableHttpError::<C::Error>::UnexpectedEndOfStream,
);
streams.spawn(Self::execute_sse_stream(
sse_stream,
sse_worker_tx.clone(),
true,
transport_task_ct.child_token(),
));
}
streams.spawn(Self::execute_sse_stream(
Self::raw_sse_to_jsonrpc(stream),
sse_worker_tx.clone(),
true,
transport_task_ct.child_token(),
));
tracing::trace!("got new sse stream after re-init");
Ok(())
}
@ -786,36 +800,12 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
Ok(())
}
Ok(StreamableHttpPostResponse::Sse(stream, ..)) => {
if let Some(session_id) = &session_id {
let sse_stream = SseAutoReconnectStream::new(
stream,
StreamableHttpClientReconnect {
client: self.client.clone(),
session_id: session_id.clone(),
uri: config.uri.clone(),
auth_header: config.auth_header.clone(),
custom_headers: protocol_headers.clone(),
},
self.config.retry_config.clone(),
);
streams.spawn(Self::execute_sse_stream(
sse_stream,
sse_worker_tx.clone(),
true,
transport_task_ct.child_token(),
));
} else {
let sse_stream = SseAutoReconnectStream::never_reconnect(
stream,
StreamableHttpError::<C::Error>::UnexpectedEndOfStream,
);
streams.spawn(Self::execute_sse_stream(
sse_stream,
sse_worker_tx.clone(),
true,
transport_task_ct.child_token(),
));
}
streams.spawn(Self::execute_sse_stream(
Self::raw_sse_to_jsonrpc(stream),
sse_worker_tx.clone(),
true,
transport_task_ct.child_token(),
));
tracing::trace!("got new sse stream");
Ok(())
}

10
crates/rmcp/src/transport/streamable_http_server/session/local.rs
View file

@ -470,7 +470,7 @@ impl LocalSessionWorker {
{
OutboundChannel::RequestWise {
id: *id,
close: false,
close: true,
}
} else {
OutboundChannel::Common
@ -483,7 +483,7 @@ impl LocalSessionWorker {
{
OutboundChannel::RequestWise {
id: *id,
close: false,
close: true,
}
} else {
OutboundChannel::Common
@ -501,7 +501,11 @@ impl LocalSessionWorker {
if let Some(request_wise) = self.tx_router.get_mut(&id) {
request_wise.tx.send(message).await;
if close {
self.tx_router.remove(&id);
if let Some(channel) = self.tx_router.remove(&id) {
for resource in channel.resources {
self.resource_router.remove(&resource);
}
}
}
} else {
return Err(SessionError::ChannelClosed(Some(id)));