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(sse): loop instead of recursing when skipping SSE events (#1146)

Browse source 
`SseAutoReconnectStream::poll_next` recursed into itself for every event it
skips — control frames, data-less frames, and frames whose data fails to
deserialize — plus once more on every state transition.

The inner stream returns `Poll::Ready` for each event it can parse out of
already-buffered bytes, so a burst of skipped events has no yield point
between them. Each one adds a stack frame, and `poll_next` is a large frame.
A client connected to a server that emits a run of non-JSON `message` frames
overflows the stack and aborts the process.

Wrap the body in a `loop` and replace the four `self.poll_next(cx)` tail calls
with `continue`. `this` is re-derived from `self.as_mut().project()` at the top
of each iteration, so the borrows end cleanly per iteration; no other logic
changes. The diff is mostly the resulting re-indent — review with `?w=1`.

Adds `skipped_events_do_not_grow_the_stack`, which feeds 50,000 undeserializable
frames through the stream. Before this change it aborts with
`fatal runtime error: stack overflow` (SIGABRT); after, it passes.
This commit is contained in:
Jeremy Schoemaker 2026-08-07 09:34:18 -05:00 • committed by GitHub
commit 3c8fb2a7e4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 164 additions and 124 deletions
Download patch file Download diff file Expand all files Collapse all files

288
crates/rmcp/src/transport/common/client_side_sse.rs
View file

@ -376,145 +376,157 @@ where
mut self: Pin<&mut Self>, mut self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>, cx: &mut std::task::Context<'_>,
) -> Poll<Option<Self::Item>> { ) -> Poll<Option<Self::Item>> {
let mut this = self.as_mut().project(); loop {
// let this_state = this.state.as_mut().project() let mut this = self.as_mut().project();
let state = this.state.as_mut().project(); // let this_state = this.state.as_mut().project()
let next_state = match state { let state = this.state.as_mut().project();
SseAutoReconnectStreamStateProj::Connected { stream } => { let next_state = match state {
match ready!(stream.poll_next(cx)) { SseAutoReconnectStreamStateProj::Connected { stream } => {
Some(Ok(sse)) => { match ready!(stream.poll_next(cx)) {
if let Some(new_server_retry) = sse.retry { Some(Ok(sse)) => {
*this.server_retry_interval = if let Some(new_server_retry) = sse.retry {
Some(Duration::from_millis(new_server_retry)); *this.server_retry_interval =
} Some(Duration::from_millis(new_server_retry));
if let Some(ref event_id) = sse.id { }
*this.last_event_id = Some(event_id.clone()); if let Some(ref event_id) = sse.id {
} *this.last_event_id = Some(event_id.clone());
// Only treat blank/`message` events as JSON-RPC payloads. }
// Other control frames (endpoint, ping, etc.) are passed to // Only treat blank/`message` events as JSON-RPC payloads.
// the reconnection handler. // Other control frames (endpoint, ping, etc.) are passed to
let is_message_event = // the reconnection handler.
matches!(sse.event.as_deref(), None | Some("") | Some("message")); let is_message_event =
if !is_message_event { matches!(sse.event.as_deref(), None | Some("") | Some("message"));
match this.connector.handle_control_event(&sse) { if !is_message_event {
Ok(()) => return self.poll_next(cx), match this.connector.handle_control_event(&sse) {
Err(e) => { Ok(()) => continue,
this.state.set(SseAutoReconnectStreamState::Terminated); Err(e) => {
return Poll::Ready(Some(Err(e))); this.state.set(SseAutoReconnectStreamState::Terminated);
return Poll::Ready(Some(Err(e)));
}
} }
} }
if let Some(data) = sse.data {
match serde_json::from_str::<ServerJsonRpcMessage>(&data) {
Err(e) => {
// Downgrade to debug to avoid noisy logs when servers emit
// non-JSON payloads as message frames. Include last_event_id
// to aid troubleshooting while keeping default behaviour.
let last_id = this.last_event_id.as_deref().unwrap_or("");
tracing::debug!(last_event_id=%last_id, "failed to deserialize server message: {e}");
continue;
}
Ok(message) => {
return Poll::Ready(Some(Ok(message)));
}
};
} else {
continue;
}
} }
if let Some(data) = sse.data { Some(Err(e)) => {
match serde_json::from_str::<ServerJsonRpcMessage>(&data) { if is_event_too_large_error(&e) {
Err(e) => { this.state.set(SseAutoReconnectStreamState::Terminated);
// Downgrade to debug to avoid noisy logs when servers emit return Poll::Ready(
// non-JSON payloads as message frames. Include last_event_id this.connector.map_fatal_stream_error(e).map(Err),
// to aid troubleshooting while keeping default behaviour. );
let last_id = this.last_event_id.as_deref().unwrap_or(""); }
tracing::debug!(last_event_id=%last_id, "failed to deserialize server message: {e}"); if *this.reconnect_only_after_event_id && this.last_event_id.is_none() {
return self.poll_next(cx); this.state.set(SseAutoReconnectStreamState::Terminated);
} return Poll::Ready(
Ok(message) => { this.connector.map_fatal_stream_error(e).map(Err),
return Poll::Ready(Some(Ok(message))); );
} }
}; this.connector
} else { .handle_stream_error(&e, this.last_event_id.as_deref());
return self.poll_next(cx); let retrying = this
} .connector
} .retry_connection(this.last_event_id.as_deref());
Some(Err(e)) => { SseAutoReconnectStreamState::Retrying {
if is_event_too_large_error(&e) {
this.state.set(SseAutoReconnectStreamState::Terminated);
return Poll::Ready(this.connector.map_fatal_stream_error(e).map(Err));
}
if *this.reconnect_only_after_event_id && this.last_event_id.is_none() {
this.state.set(SseAutoReconnectStreamState::Terminated);
return Poll::Ready(this.connector.map_fatal_stream_error(e).map(Err));
}
this.connector
.handle_stream_error(&e, this.last_event_id.as_deref());
let retrying = this
.connector
.retry_connection(this.last_event_id.as_deref());
SseAutoReconnectStreamState::Retrying {
retry_times: 0,
retrying,
}
}
None => {
if *this.reconnect_only_after_event_id && this.last_event_id.is_none() {
tracing::debug!(
"sse response ended before an event ID was received; cannot resume"
);
this.state.set(SseAutoReconnectStreamState::Terminated);
return Poll::Ready(None);
}
// Per SEP-1699, a graceful stream close is
// reconnectable. If the server sent a `retry` field
// we MUST wait that long before reconnecting.
let interval = this
.server_retry_interval
.take()
.or_else(|| this.retry_policy.retry(0));
if let Some(interval) = interval {
tracing::debug!(?interval, "sse stream ended gracefully, reconnecting");
SseAutoReconnectStreamState::WaitingNextRetry {
sleep: tokio::time::sleep(interval),
retry_times: 0, retry_times: 0,
retrying,
} }
} else {
tracing::debug!("sse stream terminated, no reconnect policy");
return Poll::Ready(None);
} }
} None => {
} if *this.reconnect_only_after_event_id && this.last_event_id.is_none() {
} tracing::debug!(
SseAutoReconnectStreamStateProj::Retrying { "sse response ended before an event ID was received; cannot resume"
retry_times, );
retrying, this.state.set(SseAutoReconnectStreamState::Terminated);
} => { return Poll::Ready(None);
let retry_result = ready!(retrying.poll(cx)); }
match retry_result { // Per SEP-1699, a graceful stream close is
Ok(new_stream) => SseAutoReconnectStreamState::Connected { stream: new_stream }, // reconnectable. If the server sent a `retry` field
Err(e) => { // we MUST wait that long before reconnecting.
tracing::debug!("retry sse stream error: {e}");
*retry_times += 1;
if let Some(interval) = this.retry_policy.retry(*retry_times) {
let interval = this let interval = this
.server_retry_interval .server_retry_interval
.map(|server_retry_interval| server_retry_interval.max(interval)) .take()
.unwrap_or(interval); .or_else(|| this.retry_policy.retry(0));
let sleep = tokio::time::sleep(interval); if let Some(interval) = interval {
SseAutoReconnectStreamState::WaitingNextRetry { tracing::debug!(
sleep, ?interval,
retry_times: *retry_times, "sse stream ended gracefully, reconnecting"
);
SseAutoReconnectStreamState::WaitingNextRetry {
sleep: tokio::time::sleep(interval),
retry_times: 0,
}
} else {
tracing::debug!("sse stream terminated, no reconnect policy");
return Poll::Ready(None);
} }
} else {
tracing::error!("sse stream error: {e}, max retry times reached");
this.state.set(SseAutoReconnectStreamState::Terminated);
return Poll::Ready(Some(Err(e)));
} }
} }
} }
} SseAutoReconnectStreamStateProj::Retrying {
SseAutoReconnectStreamStateProj::WaitingNextRetry { sleep, retry_times } => {
ready!(sleep.poll(cx));
let retrying = this
.connector
.retry_connection(this.last_event_id.as_deref());
let retry_times = *retry_times;
SseAutoReconnectStreamState::Retrying {
retry_times, retry_times,
retrying, retrying,
} => {
let retry_result = ready!(retrying.poll(cx));
match retry_result {
Ok(new_stream) => {
SseAutoReconnectStreamState::Connected { stream: new_stream }
}
Err(e) => {
tracing::debug!("retry sse stream error: {e}");
*retry_times += 1;
if let Some(interval) = this.retry_policy.retry(*retry_times) {
let interval = this
.server_retry_interval
.map(|server_retry_interval| {
server_retry_interval.max(interval)
})
.unwrap_or(interval);
let sleep = tokio::time::sleep(interval);
SseAutoReconnectStreamState::WaitingNextRetry {
sleep,
retry_times: *retry_times,
}
} else {
tracing::error!("sse stream error: {e}, max retry times reached");
this.state.set(SseAutoReconnectStreamState::Terminated);
return Poll::Ready(Some(Err(e)));
}
}
}
} }
} SseAutoReconnectStreamStateProj::WaitingNextRetry { sleep, retry_times } => {
SseAutoReconnectStreamStateProj::Terminated => { ready!(sleep.poll(cx));
return Poll::Ready(None); let retrying = this
} .connector
}; .retry_connection(this.last_event_id.as_deref());
// update the state let retry_times = *retry_times;
this.state.set(next_state); SseAutoReconnectStreamState::Retrying {
self.poll_next(cx) retry_times,
retrying,
}
}
SseAutoReconnectStreamStateProj::Terminated => {
return Poll::Ready(None);
}
};
// update the state
this.state.set(next_state);
}
} }
} }
@ -716,6 +728,34 @@ mod tests {
); );
} }
#[tokio::test]
async fn skipped_events_do_not_grow_the_stack() {
// Every frame here fails to deserialize, so each one takes the "skip this
// event" path. The inner stream is always ready, so nothing yields in
// between: when that path recursed into poll_next instead of looping, this
// overflowed the stack and aborted the process.
const SKIPPED_EVENTS: usize = 50_000;
let mut payload = Vec::with_capacity(SKIPPED_EVENTS * 9);
for _ in 0..SKIPPED_EVENTS {
payload.extend_from_slice(b"data: x\n\n");
}
let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from(payload))]);
let attempts = Arc::new(AtomicUsize::new(0));
let connector = CountingReconnect {
attempts: attempts.clone(),
};
let stream = SseAutoReconnectStream::new(
bounded_sse_stream(source, 1024),
connector,
Arc::new(NeverRetry),
);
let mut stream = std::pin::pin!(stream);
assert!(stream.next().await.is_none());
assert_eq!(attempts.load(Ordering::Relaxed), 0);
}
#[tokio::test] #[tokio::test]
async fn response_without_event_id_does_not_reconnect() { async fn response_without_event_id_does_not_reconnect() {
let attempts = Arc::new(AtomicUsize::new(0)); let attempts = Arc::new(AtomicUsize::new(0));