portugalfuturista/rust-sdk
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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
parent 6603c1ff15
commit 3e56d52764
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 467 additions and 89 deletions
Show stats 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,

340
crates/rmcp/tests/test_streamable_http_priming.rs
View file

@ -2,7 +2,8 @@
use std::time::Duration;
use rmcp::transport::streamable_http_server::{
StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
StreamableHttpServerConfig, StreamableHttpService,
session::{SessionId, local::LocalSessionManager},
};
use tokio_util::sync::CancellationToken;
@ -54,7 +55,7 @@ async fn test_priming_on_stream_start() -> anyhow::Result<()> {
let events: Vec<&str> = body.split("\n\n").filter(|e| !e.is_empty()).collect();
assert!(events.len() >= 2);
// Verify priming event (first event)
// Verify priming event (first event) — initialize uses "0" (no http_request_id)
let priming_event = events[0];
assert!(priming_event.contains("id: 0"));
assert!(priming_event.contains("retry: 3000"));
@ -71,6 +72,341 @@ async fn test_priming_on_stream_start() -> anyhow::Result<()> {
Ok(())
}
#[tokio::test]
async fn test_request_wise_priming_includes_http_request_id() -> anyhow::Result<()> {
let ct = CancellationToken::new();
let service: StreamableHttpService<Calculator, LocalSessionManager> =
StreamableHttpService::new(
|| Ok(Calculator::new()),
Default::default(),
StreamableHttpServerConfig::default()
.with_sse_keep_alive(None)
.with_cancellation_token(ct.child_token()),
);
let router = axum::Router::new().nest_service("/mcp", service);
let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let addr = tcp_listener.local_addr()?;
let handle = tokio::spawn({
let ct = ct.clone();
async move {
let _ = axum::serve(tcp_listener, router)
.with_graceful_shutdown(async move { ct.cancelled_owned().await })
.await;
}
});
let client = reqwest::Client::new();
// Initialize the session
let response = client
.post(format!("http://{addr}/mcp"))
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.body(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#)
.send()
.await?;
assert_eq!(response.status(), 200);
let session_id: SessionId = response.headers()["mcp-session-id"].to_str()?.into();
// Send notifications/initialized
let status = client
.post(format!("http://{addr}/mcp"))
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.header("mcp-session-id", session_id.to_string())
.header("Mcp-Protocol-Version", "2025-06-18")
.body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
.send()
.await?
.status();
assert_eq!(status, 202);
// First tool call — should get http_request_id 0
let body = client
.post(format!("http://{addr}/mcp"))
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.header("mcp-session-id", session_id.to_string())
.header("Mcp-Protocol-Version", "2025-06-18")
.body(r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"sum","arguments":{"a":1,"b":2}}}"#)
.send()
.await?
.text()
.await?;
let events: Vec<&str> = body.split("\n\n").filter(|e| !e.is_empty()).collect();
assert!(
events.len() >= 2,
"expected priming + response, got: {body}"
);
// Priming event should encode the http_request_id (0)
let priming = events[0];
assert!(
priming.contains("id: 0/0"),
"first request priming should be 0/0, got: {priming}"
);
assert!(priming.contains("retry: 3000"));
// Response event should use index 1 (since priming occupies index 0)
let response_event = events[1];
assert!(
response_event.contains("id: 1/0"),
"first response event id should be 1/0, got: {response_event}"
);
assert!(response_event.contains(r#""id":2"#));
// Second tool call — should get http_request_id 1
let body = client
.post(format!("http://{addr}/mcp"))
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.header("mcp-session-id", session_id.to_string())
.header("Mcp-Protocol-Version", "2025-06-18")
.body(r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"sum","arguments":{"a":3,"b":4}}}"#)
.send()
.await?
.text()
.await?;
let events: Vec<&str> = body.split("\n\n").filter(|e| !e.is_empty()).collect();
assert!(
events.len() >= 2,
"expected priming + response, got: {body}"
);
let priming = events[0];
assert!(
priming.contains("id: 0/1"),
"second request priming should be 0/1, got: {priming}"
);
let response_event = events[1];
assert!(
response_event.contains("id: 1/1"),
"second response event id should be 1/1, got: {response_event}"
);
assert!(response_event.contains(r#""id":3"#));
ct.cancel();
handle.await?;
Ok(())
}
#[tokio::test]
async fn test_resume_after_request_wise_channel_completed() -> anyhow::Result<()> {
let ct = CancellationToken::new();
let service: StreamableHttpService<Calculator, LocalSessionManager> =
StreamableHttpService::new(
|| Ok(Calculator::new()),
Default::default(),
StreamableHttpServerConfig::default()
.with_sse_keep_alive(None)
.with_cancellation_token(ct.child_token()),
);
let router = axum::Router::new().nest_service("/mcp", service);
let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let addr = tcp_listener.local_addr()?;
let handle = tokio::spawn({
let ct = ct.clone();
async move {
let _ = axum::serve(tcp_listener, router)
.with_graceful_shutdown(async move { ct.cancelled_owned().await })
.await;
}
});
let client = reqwest::Client::new();
// Initialize session
let response = client
.post(format!("http://{addr}/mcp"))
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.body(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#)
.send()
.await?;
assert_eq!(response.status(), 200);
let session_id: SessionId = response.headers()["mcp-session-id"].to_str()?.into();
// Complete handshake
let status = client
.post(format!("http://{addr}/mcp"))
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.header("mcp-session-id", session_id.to_string())
.header("Mcp-Protocol-Version", "2025-06-18")
.body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
.send()
.await?
.status();
assert_eq!(status, 202);
// Call a tool and consume the full response (channel completes)
let body = client
.post(format!("http://{addr}/mcp"))
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.header("mcp-session-id", session_id.to_string())
.header("Mcp-Protocol-Version", "2025-06-18")
.body(r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"sum","arguments":{"a":1,"b":2}}}"#)
.send()
.await?
.text()
.await?;
let events: Vec<&str> = body.split("\n\n").filter(|e| !e.is_empty()).collect();
assert!(
events.len() >= 2,
"expected priming + response, got: {body}"
);
assert!(events[0].contains("id: 0/0"));
assert!(events[1].contains(r#""id":2"#));
// Resume with Last-Event-ID after the channel has completed.
// The server returns 200 — either with replayed cached events
// (if the channel is still retained) or an empty stream (if the
// session worker hasn't processed the completion yet).
let resume = client
.get(format!("http://{addr}/mcp"))
.header("Accept", "text/event-stream")
.header("mcp-session-id", session_id.to_string())
.header("Mcp-Protocol-Version", "2025-06-18")
.header("last-event-id", "0/0")
.send()
.await?;
assert_eq!(resume.status(), 200);
let resume_body = resume.text().await?;
// The stream should complete (not hang), regardless of whether
// it contains replayed events or is empty.
assert!(
!resume_body.contains("standalone"),
"should not receive events from a different stream"
);
ct.cancel();
handle.await?;
Ok(())
}
#[tokio::test]
async fn test_completed_cache_ttl_eviction() -> anyhow::Result<()> {
use std::sync::Arc;
let ct = CancellationToken::new();
let mut session_manager = LocalSessionManager::default();
session_manager.session_config.completed_cache_ttl = Duration::from_millis(200);
let session_manager = Arc::new(session_manager);
let service = StreamableHttpService::new(
|| Ok(Calculator::new()),
session_manager.clone(),
StreamableHttpServerConfig::default()
.with_sse_keep_alive(None)
.with_cancellation_token(ct.child_token()),
);
let router = axum::Router::new().nest_service("/mcp", service);
let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let addr = tcp_listener.local_addr()?;
let handle = tokio::spawn({
let ct = ct.clone();
async move {
let _ = axum::serve(tcp_listener, router)
.with_graceful_shutdown(async move { ct.cancelled_owned().await })
.await;
}
});
let client = reqwest::Client::new();
// Initialize session
let response = client
.post(format!("http://{addr}/mcp"))
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.body(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#)
.send()
.await?;
assert_eq!(response.status(), 200);
let session_id: SessionId = response.headers()["mcp-session-id"].to_str()?.into();
// Complete handshake
client
.post(format!("http://{addr}/mcp"))
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.header("mcp-session-id", session_id.to_string())
.header("Mcp-Protocol-Version", "2025-06-18")
.body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
.send()
.await?;
// Call a tool and consume the response (channel completes)
let body = client
.post(format!("http://{addr}/mcp"))
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.header("mcp-session-id", session_id.to_string())
.header("Mcp-Protocol-Version", "2025-06-18")
.body(r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"sum","arguments":{"a":1,"b":2}}}"#)
.send()
.await?
.text()
.await?;
assert!(body.contains(r#""id":2"#));
// Wait for TTL to expire (200ms) plus margin
tokio::time::sleep(Duration::from_millis(400)).await;
// Send a notification to trigger an event loop iteration (runs eviction)
client
.post(format!("http://{addr}/mcp"))
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.header("mcp-session-id", session_id.to_string())
.header("Mcp-Protocol-Version", "2025-06-18")
.body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
.send()
.await?;
// Small delay to ensure the eviction ran
tokio::time::sleep(Duration::from_millis(50)).await;
// Resume after TTL — channel should be evicted. The server returns
// 200 with an empty stream (no events from a different stream).
let resume = client
.get(format!("http://{addr}/mcp"))
.header("Accept", "text/event-stream")
.header("mcp-session-id", session_id.to_string())
.header("Mcp-Protocol-Version", "2025-06-18")
.header("last-event-id", "0/0")
.send()
.await?;
assert_eq!(resume.status(), 200);
let body = resume.text().await?;
assert!(
!body.contains(r#""id":2"#),
"should NOT contain the old tool response after eviction, got: {body}"
);
ct.cancel();
handle.await?;
Ok(())
}
#[tokio::test]
async fn test_priming_on_stream_close() -> anyhow::Result<()> {
use std::sync::Arc;