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

feat: add configuration for transparent session re-init (#760)

Browse source 
* feat: add configuration for transparent session re-init

* fix: in ci revert running tests without local until all tests pass

* fix: pr comments

* fix: documentation
This commit is contained in:
Guy Lichtman 2026-03-22 21:05:40 +02:00 • committed by GitHub
commit 3ea8c3c555
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 203 additions and 95 deletions
Download patch file Download diff file Expand all files Collapse all files

12
crates/rmcp/Cargo.toml
View file

@ -265,3 +265,15 @@ path = "tests/test_sse_concurrent_streams.rs"
name = "test_client_credentials" name = "test_client_credentials"
required-features = ["auth"] required-features = ["auth"]
path = "tests/test_client_credentials.rs" path = "tests/test_client_credentials.rs"
[[test]]
name = "test_streamable_http_stale_session"
required-features = [
"server",
"client",
"transport-streamable-http-server",
"transport-streamable-http-client",
"transport-streamable-http-client-reqwest"
]
path = "tests/test_streamable_http_stale_session.rs"

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

@ -600,48 +600,51 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
.await; .await;
let send_result = match response { let send_result = match response {
Err(StreamableHttpError::SessionExpired) => { Err(StreamableHttpError::SessionExpired) => {
// The server discarded the session (HTTP 404). Perform a if !config.reinit_on_expired_session {
// fresh handshake once and replay the original message. Err(StreamableHttpError::SessionExpired)
tracing::info!( } else {
"session expired (HTTP 404), attempting transparent re-initialization" // The server discarded the session (HTTP 404). Perform a
); // fresh handshake once and replay the original message.
match Self::perform_reinitialization( tracing::info!(
self.client.clone(), "session expired (HTTP 404), attempting transparent re-initialization"
saved_init_request.clone(), );
config.uri.clone(), match Self::perform_reinitialization(
config.auth_header.clone(), self.client.clone(),
config.custom_headers.clone(), saved_init_request.clone(),
) config.uri.clone(),
.await config.auth_header.clone(),
{ config.custom_headers.clone(),
Ok((new_session_id, new_protocol_headers)) => { )
// Old streams hold the stale session ID; abort them .await
// so the new standalone SSE stream takes over. {
streams.abort_all(); Ok((new_session_id, new_protocol_headers)) => {
// Old streams hold the stale session ID; abort them
// so the new standalone SSE stream takes over.
streams.abort_all();
session_id = new_session_id; session_id = new_session_id;
protocol_headers = new_protocol_headers; protocol_headers = new_protocol_headers;
session_cleanup_info = session_cleanup_info =
session_id.as_ref().map(|sid| SessionCleanupInfo { session_id.as_ref().map(|sid| SessionCleanupInfo {
client: self.client.clone(), client: self.client.clone(),
uri: config.uri.clone(), uri: config.uri.clone(),
session_id: sid.clone(), session_id: sid.clone(),
auth_header: config.auth_header.clone(), auth_header: config.auth_header.clone(),
protocol_headers: protocol_headers.clone(), protocol_headers: protocol_headers.clone(),
}); });
if let Some(new_sid) = &session_id { if let Some(new_sid) = &session_id {
let client = self.client.clone(); let client = self.client.clone();
let uri = config.uri.clone(); let uri = config.uri.clone();
let new_sid = new_sid.clone(); let new_sid = new_sid.clone();
let auth_header = config.auth_header.clone(); let auth_header = config.auth_header.clone();
let retry_config = self.config.retry_config.clone(); let retry_config = self.config.retry_config.clone();
let sse_tx = sse_worker_tx.clone(); let sse_tx = sse_worker_tx.clone();
let task_ct = transport_task_ct.clone(); let task_ct = transport_task_ct.clone();
let config_uri = config.uri.clone(); let config_uri = config.uri.clone();
let config_auth = config.auth_header.clone(); let config_auth = config.auth_header.clone();
let spawn_headers = protocol_headers.clone(); let spawn_headers = protocol_headers.clone();
streams.spawn(async move { streams.spawn(async move {
match client match client
.get_stream( .get_stream(
uri, uri,
@ -686,69 +689,71 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
} }
} }
}); });
} }
let retry_response = self let retry_response = self
.client .client
.post_message( .post_message(
config.uri.clone(), config.uri.clone(),
message, message,
session_id.clone(), session_id.clone(),
config.auth_header.clone(), config.auth_header.clone(),
protocol_headers.clone(), protocol_headers.clone(),
) )
.await; .await;
match retry_response { match retry_response {
Err(e) => Err(e), Err(e) => Err(e),
Ok(StreamableHttpPostResponse::Accepted) => { Ok(StreamableHttpPostResponse::Accepted) => {
tracing::trace!( tracing::trace!(
"client message accepted after re-init" "client message accepted after re-init"
);
Ok(())
}
Ok(StreamableHttpPostResponse::Json(msg, ..)) => {
context.send_to_handler(msg).await?;
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( Ok(())
sse_stream, }
sse_worker_tx.clone(), Ok(StreamableHttpPostResponse::Json(msg, ..)) => {
true, context.send_to_handler(msg).await?;
transport_task_ct.child_token(), Ok(())
)); }
} else { Ok(StreamableHttpPostResponse::Sse(stream, ..)) => {
let 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( SseAutoReconnectStream::never_reconnect(
stream, stream,
StreamableHttpError::<C::Error>::UnexpectedEndOfStream, StreamableHttpError::<C::Error>::UnexpectedEndOfStream,
); );
streams.spawn(Self::execute_sse_stream( streams.spawn(Self::execute_sse_stream(
sse_stream, sse_stream,
sse_worker_tx.clone(), sse_worker_tx.clone(),
true, true,
transport_task_ct.child_token(), transport_task_ct.child_token(),
)); ));
}
tracing::trace!("got new sse stream after re-init");
Ok(())
} }
tracing::trace!("got new sse stream after re-init");
Ok(())
} }
} }
Err(reinit_err) => Err(reinit_err),
} }
Err(reinit_err) => Err(reinit_err), } // else enable_reinit_on_expired_session
}
} }
Err(e) => Err(e), Err(e) => Err(e),
Ok(StreamableHttpPostResponse::Accepted) => { Ok(StreamableHttpPostResponse::Accepted) => {
@ -1051,6 +1056,16 @@ pub struct StreamableHttpClientTransportConfig {
pub auth_header: Option<String>, pub auth_header: Option<String>,
/// Custom HTTP headers to include with every request /// Custom HTTP headers to include with every request
pub custom_headers: HashMap<HeaderName, HeaderValue>, pub custom_headers: HashMap<HeaderName, HeaderValue>,
/// Enables transparent recovery when the server reports an expired session (`HTTP 404`).
///
/// When enabled, the transport performs one automatic recovery attempt:
/// 1. Replays the original `initialize` handshake to create a new session.
/// 2. Re-establishes streaming state for that session.
/// 3. Retries the in-flight request that failed with `SessionExpired`.
///
/// This recovery is best-effort and bounded to a single attempt. If recovery fails,
/// the original failure path is preserved and the error is returned to the caller.
pub reinit_on_expired_session: bool,
} }
impl StreamableHttpClientTransportConfig { impl StreamableHttpClientTransportConfig {
@ -1098,6 +1113,19 @@ impl StreamableHttpClientTransportConfig {
self.custom_headers = custom_headers; self.custom_headers = custom_headers;
self self
} }
/// Set whether the transport should attempt transparent re-initialization on session expiration
/// See [`Self::reinit_on_expired_session`] for details.
/// # Example
/// ```rust,no_run
/// use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
/// let config = StreamableHttpClientTransportConfig::with_uri("http://localhost:8000")
/// .reinit_on_expired_session(true);
/// ```
pub fn reinit_on_expired_session(mut self, enable: bool) -> Self {
self.reinit_on_expired_session = enable;
self
}
} }
impl Default for StreamableHttpClientTransportConfig { impl Default for StreamableHttpClientTransportConfig {
@ -1109,6 +1137,7 @@ impl Default for StreamableHttpClientTransportConfig {
allow_stateless: true, allow_stateless: true,
auth_header: None, auth_header: None,
custom_headers: HashMap::new(), custom_headers: HashMap::new(),
reinit_on_expired_session: true,
} }
} }
} }

71
crates/rmcp/tests/test_streamable_http_stale_session.rs
View file

@ -8,7 +8,7 @@
use std::{collections::HashMap, sync::Arc}; use std::{collections::HashMap, sync::Arc};
use rmcp::{ use rmcp::{
ServiceExt, ServiceError, ServiceExt,
model::{ClientJsonRpcMessage, ClientRequest, PingRequest, RequestId}, model::{ClientJsonRpcMessage, ClientRequest, PingRequest, RequestId},
transport::{ transport::{
StreamableHttpClientTransport, StreamableHttpClientTransport,
@ -126,7 +126,8 @@ async fn test_transparent_reinitialization_on_session_expiry() -> anyhow::Result
// Connect a full client transport (this performs initialize + notifications/initialized) // Connect a full client transport (this performs initialize + notifications/initialized)
let transport = StreamableHttpClientTransport::from_config( let transport = StreamableHttpClientTransport::from_config(
StreamableHttpClientTransportConfig::with_uri(format!("http://{addr}/mcp")), StreamableHttpClientTransportConfig::with_uri(format!("http://{addr}/mcp"))
.reinit_on_expired_session(true),
); );
let client = ().serve(transport).await?; let client = ().serve(transport).await?;
@ -171,3 +172,69 @@ async fn test_transparent_reinitialization_on_session_expiry() -> anyhow::Result
Ok(()) Ok(())
} }
/// Verify that when `reinit_on_expired_session` is false and the server loses the session,
/// the client receives a `SessionExpired` transport error instead of retrying.
#[tokio::test]
async fn test_session_expired_error_when_reinit_disabled() -> anyhow::Result<()> {
let ct = CancellationToken::new();
let session_manager = Arc::new(LocalSessionManager::default());
let service = StreamableHttpService::new(
|| Ok(Calculator::new()),
session_manager.clone(),
StreamableHttpServerConfig {
stateful_mode: true,
sse_keep_alive: None,
cancellation_token: ct.child_token(),
..Default::default()
},
);
let router = axum::Router::new().nest_service("/mcp", service);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let addr = listener.local_addr()?;
let server_handle = tokio::spawn({
let ct = ct.clone();
async move {
let _ = axum::serve(listener, router)
.with_graceful_shutdown(async move { ct.cancelled_owned().await })
.await;
}
});
let transport = StreamableHttpClientTransport::from_config(
StreamableHttpClientTransportConfig::with_uri(format!("http://{addr}/mcp"))
.reinit_on_expired_session(false),
);
let client = ().serve(transport).await?;
// Verify the session is established
let _resources = client.list_all_resources().await?;
// Force session expiry by removing all sessions from the server-side manager
{
let mut sessions = session_manager.sessions.write().await;
sessions.clear();
}
// This call should fail with a SessionExpired transport error
let result = client.list_all_resources().await;
match result {
Err(ServiceError::TransportSend(ref dyn_err)) => {
let err_msg = format!("{dyn_err}");
assert!(
err_msg.contains("Session expired"),
"expected 'Session expired' in error message, got: {err_msg}"
);
}
other => panic!("expected TransportSend(SessionExpired), got: {other:?}"),
}
let _ = client.cancel().await;
ct.cancel();
server_handle.await?;
Ok(())
}