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: include http_request_id in request-wise priming event IDs (#799)

Browse source 
* fix: include http_request_id in request-wise priming event IDs

* refactor: use Option::into_iter and usize::from for priming

* fix: retain event cache for completed request-wise channels

* fix: track completed_at for cache eviction and resume

* fix: log resume failures at warn level

* test: add completed_cache_ttl eviction test

* fix: return empty stream on failed resume

* test: add resume after completion test
This commit is contained in:
Dale Seo 2026-04-16 12:02:35 -04:00 • committed by GitHub
commit 3e56d52764
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 467 additions and 89 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1,10 +1,10 @@
use std::{
collections::{HashMap, HashSet, VecDeque},
num::ParseIntError,
time::Duration,
time::{Duration, Instant},
};
use futures::Stream;
use futures::{Stream, StreamExt};
use thiserror::Error;
use tokio::sync::{
mpsc::{Receiver, Sender},
@ -86,10 +86,17 @@ impl SessionManager for LocalSessionManager {
.get(id)
.ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
let receiver = handle.establish_request_wise_channel().await?;
handle
.push_message(message, receiver.http_request_id)
.await?;
Ok(ReceiverStream::new(receiver.inner))
let http_request_id = receiver.http_request_id;
handle.push_message(message, http_request_id).await?;
let priming = self.session_config.sse_retry.map(|retry| {
let event_id = match http_request_id {
Some(id) => format!("0/{id}"),
None => "0".into(),
};
ServerSseMessage::priming(event_id, retry)
});
Ok(futures::stream::iter(priming).chain(ReceiverStream::new(receiver.inner)))
}
async fn create_standalone_stream(
@ -188,23 +195,29 @@ struct CachedTx {
cache: VecDeque<ServerSseMessage>,
http_request_id: Option<HttpRequestId>,
capacity: usize,
starting_index: usize,
}
impl CachedTx {
fn new(tx: Sender<ServerSseMessage>, http_request_id: Option<HttpRequestId>) -> Self {
fn new(
tx: Sender<ServerSseMessage>,
http_request_id: Option<HttpRequestId>,
starting_index: usize,
) -> Self {
Self {
cache: VecDeque::with_capacity(tx.capacity()),
capacity: tx.capacity(),
tx,
http_request_id,
starting_index,
}
}
fn new_common(tx: Sender<ServerSseMessage>) -> Self {
Self::new(tx, None)
Self::new(tx, None, 0)
}
fn next_event_id(&self) -> EventId {
let index = self.cache.back().map_or(0, |m| {
let index = self.cache.back().map_or(self.starting_index, |m| {
m.event_id
.as_deref()
.unwrap_or_default()
@ -272,6 +285,7 @@ impl CachedTx {
struct HttpRequestWise {
resources: HashSet<ResourceKey>,
tx: CachedTx,
completed_at: Option<Instant>,
}
type HttpRequestId = u64;
@ -342,23 +356,27 @@ pub struct StreamableHttpMessageReceiver {
impl LocalSessionWorker {
fn unregister_resource(&mut self, resource: &ResourceKey) {
if let Some(http_request_id) = self.resource_router.remove(resource) {
tracing::trace!(?resource, http_request_id, "unregister resource");
if let Some(channel) = self.tx_router.get_mut(&http_request_id) {
// It's okey to do so, since we don't handle batch json rpc request anymore
// and this can be refactored after the batch request is removed in the coming version.
if channel.resources.is_empty() || matches!(resource, ResourceKey::McpRequestId(_))
{
tracing::debug!(http_request_id, "close http request wise channel");
if let Some(channel) = self.tx_router.remove(&http_request_id) {
for resource in channel.resources {
self.resource_router.remove(&resource);
}
}
}
} else {
tracing::warn!(http_request_id, "http request wise channel not found");
}
let Some(http_request_id) = self.resource_router.remove(resource) else {
return;
};
tracing::trace!(?resource, http_request_id, "unregister resource");
let Some(channel) = self.tx_router.get_mut(&http_request_id) else {
tracing::warn!(http_request_id, "http request wise channel not found");
return;
};
if !channel.resources.is_empty() && !matches!(resource, ResourceKey::McpRequestId(_)) {
return;
}
tracing::debug!(http_request_id, "close http request wise channel");
let resources: Vec<_> = channel.resources.drain().collect();
channel.completed_at = Some(Instant::now());
// Close the sender so the client's SSE stream ends,
// but keep the entry so the cache is available for
// late resume requests.
let (closed_tx, _) = tokio::sync::mpsc::channel(1);
channel.tx.tx = closed_tx;
for resource in resources {
self.resource_router.remove(&resource);
}
}
fn register_resource(&mut self, resource: ResourceKey, http_request_id: HttpRequestId) {
@ -395,6 +413,11 @@ impl LocalSessionWorker {
self.unregister_resource(&resource);
}
}
fn evict_expired_channels(&mut self) {
let ttl = self.session_config.completed_cache_ttl;
self.tx_router
.retain(|_, rw| rw.completed_at.is_none_or(|at| at.elapsed() < ttl));
}
fn next_http_request_id(&mut self) -> HttpRequestId {
let id = self.next_http_request_id;
self.next_http_request_id = self.next_http_request_id.wrapping_add(1);
@ -405,11 +428,13 @@ impl LocalSessionWorker {
) -> Result<StreamableHttpMessageReceiver, SessionError> {
let http_request_id = self.next_http_request_id();
let (tx, rx) = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
let starting_index = usize::from(self.session_config.sse_retry.is_some());
self.tx_router.insert(
http_request_id,
HttpRequestWise {
resources: Default::default(),
tx: CachedTx::new(tx, Some(http_request_id)),
tx: CachedTx::new(tx, Some(http_request_id), starting_index),
completed_at: None,
},
);
tracing::debug!(http_request_id, "establish new request wise channel");
@ -524,28 +549,25 @@ impl LocalSessionWorker {
match last_event_id.http_request_id {
Some(http_request_id) => {
if let Some(request_wise) = self.tx_router.get_mut(&http_request_id) {
// Resume existing request-wise channel
let channel = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
let (tx, rx) = channel;
request_wise.tx.tx = tx;
let index = last_event_id.index;
// sync messages after index
request_wise.tx.sync(index).await?;
Ok(StreamableHttpMessageReceiver {
http_request_id: Some(http_request_id),
inner: rx,
})
} else {
// Request-wise channel completed (POST response already delivered).
// The client's EventSource is reconnecting after the POST SSE stream
// ended. Fall through to common channel handling below.
tracing::debug!(
http_request_id,
"Request-wise channel completed, falling back to common channel"
);
self.resume_or_shadow_common(last_event_id.index).await
let request_wise = self
.tx_router
.get_mut(&http_request_id)
.ok_or(SessionError::ChannelClosed(Some(http_request_id)))?;
let is_completed = request_wise.completed_at.is_some();
let (tx, rx) = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
request_wise.tx.tx = tx;
let index = last_event_id.index;
request_wise.tx.sync(index).await?;
if is_completed {
// Drop the sender after replaying so the stream ends
// instead of hanging indefinitely.
let (closed_tx, _) = tokio::sync::mpsc::channel(1);
request_wise.tx.tx = closed_tx;
}
Ok(StreamableHttpMessageReceiver {
http_request_id: Some(http_request_id),
inner: rx,
})
}
None => self.resume_or_shadow_common(last_event_id.index).await,
}
@ -955,6 +977,7 @@ impl Worker for LocalSessionWorker {
let ct = context.cancellation_token.clone();
let keep_alive = self.session_config.keep_alive.unwrap_or(Duration::MAX);
loop {
self.evict_expired_channels();
let keep_alive_timeout = tokio::time::sleep(keep_alive);
let event = tokio::select! {
event = self.event_rx.recv() => {
@ -1076,11 +1099,22 @@ pub struct SessionConfig {
/// Defaults to 5 minutes. Set to `None` to disable (not recommended
/// for long-running servers behind proxies).
pub keep_alive: Option<Duration>,
/// SSE retry interval for priming events on request-wise streams.
/// When set, the session layer prepends a priming event with the correct
/// stream-identifying event ID to each request-wise SSE stream.
/// Default is 3 seconds, matching `StreamableHttpServerConfig::default()`.
pub sse_retry: Option<Duration>,
/// How long to retain completed request-wise channel caches for late
/// resume requests. After this duration, completed entries are evicted
/// and resume will return an error. Default is 60 seconds.
pub completed_cache_ttl: Duration,
}
impl SessionConfig {
pub const DEFAULT_CHANNEL_CAPACITY: usize = 16;
pub const DEFAULT_KEEP_ALIVE: Duration = Duration::from_secs(300);
pub const DEFAULT_SSE_RETRY: Duration = Duration::from_secs(3);
pub const DEFAULT_COMPLETED_CACHE_TTL: Duration = Duration::from_secs(60);
}
impl Default for SessionConfig {
@ -1088,6 +1122,8 @@ impl Default for SessionConfig {
Self {
channel_capacity: Self::DEFAULT_CHANNEL_CAPACITY,
keep_alive: Some(Self::DEFAULT_KEEP_ALIVE),
sse_retry: Some(Self::DEFAULT_SSE_RETRY),
completed_cache_ttl: Self::DEFAULT_COMPLETED_CACHE_TTL,
}
}
}

84
crates/rmcp/src/transport/streamable_http_server/tower.rs
View file

@ -478,40 +478,52 @@ where
.and_then(|v| v.to_str().ok())
.map(|s| s.to_owned());
if let Some(last_event_id) = last_event_id {
// check if session has this event id
let stream = self
match self
.session_manager
.resume(&session_id, last_event_id)
.await
.map_err(internal_error_response("resume session"))?;
// Resume doesn't need priming - client already has the event ID
Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
} else {
// create standalone stream
let stream = self
.session_manager
.create_standalone_stream(&session_id)
.await
.map_err(internal_error_response("create standalone stream"))?;
// Prepend priming event if sse_retry configured
let stream = if let Some(retry) = self.config.sse_retry {
let priming = ServerSseMessage::priming("0", retry);
futures::stream::once(async move { priming })
.chain(stream)
.left_stream()
} else {
stream.right_stream()
};
Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
{
Ok(stream) => {
return Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
));
}
Err(e) => {
// Return 200 with an immediately-closed empty stream.
// Returning an HTTP error would cause EventSource to retry
// with the same Last-Event-ID in an infinite loop. An empty
// 200 cleanly terminates the EventSource without delivering
// events from a different stream.
tracing::warn!("Resume failed ({e}), returning empty stream");
return Ok(sse_stream_response(
futures::stream::empty(),
None,
self.config.cancellation_token.child_token(),
));
}
}
}
// No Last-Event-ID — create standalone stream
let stream = self
.session_manager
.create_standalone_stream(&session_id)
.await
.map_err(internal_error_response("create standalone stream"))?;
let stream = if let Some(retry) = self.config.sse_retry {
let priming = ServerSseMessage::priming("0", retry);
futures::stream::once(async move { priming })
.chain(stream)
.left_stream()
} else {
stream.right_stream()
};
Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
}
async fn handle_post<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
@ -598,20 +610,14 @@ where
match message {
ClientJsonRpcMessage::Request(_) => {
// Priming for request-wise streams is handled by the
// session layer (SessionManager::create_stream) which
// has access to the http_request_id for correct event IDs.
let stream = self
.session_manager
.create_stream(&session_id, message)
.await
.map_err(internal_error_response("get session"))?;
// Prepend priming event if sse_retry configured
let stream = if let Some(retry) = self.config.sse_retry {
let priming = ServerSseMessage::priming("0", retry);
futures::stream::once(async move { priming })
.chain(stream)
.left_stream()
} else {
stream.right_stream()
};
Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,