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

feat: transparent session re-init on HTTP 404 (#743)

Browse source
This commit is contained in:
Dale Seo 2026-03-11 10:27:00 -04:00 committed by GitHub
parent 5322430772
commit 27b00967f1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 321 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -144,6 +144,7 @@ impl StreamableHttpClient for reqwest::Client {
}
request = apply_custom_headers(request, custom_headers)?;
let session_was_attached = session_id.is_some();
if let Some(session_id) = session_id {
request = request.header(HEADER_SESSION_ID, session_id.as_ref());
}
@ -186,6 +187,9 @@ impl StreamableHttpClient for reqwest::Client {
) {
return Ok(StreamableHttpPostResponse::Accepted);
}
if status == reqwest::StatusCode::NOT_FOUND && session_was_attached {
return Err(StreamableHttpError::SessionExpired);
}
if !status.is_success() {
let body = response
.text()

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

@ -11,7 +11,10 @@ use tracing::debug;
use super::common::client_side_sse::{ExponentialBackoff, SseRetryPolicy, SseStreamReconnect};
use crate::{
RoleClient,
model::{ClientJsonRpcMessage, ServerJsonRpcMessage, ServerResult},
model::{
ClientJsonRpcMessage, ClientNotification, InitializedNotification, ServerJsonRpcMessage,
ServerResult,
},
transport::{
common::client_side_sse::SseAutoReconnectStream,
worker::{Worker, WorkerQuitReason, WorkerSendRequest, WorkerTransport},
@ -79,6 +82,8 @@ pub enum StreamableHttpError<E: std::error::Error + Send + Sync + 'static> {
InsufficientScope(InsufficientScopeError),
#[error("Header name '{0}' is reserved and conflicts with default headers")]
ReservedHeaderConflict(String),
#[error("Session expired (HTTP 404)")]
SessionExpired,
}
#[derive(Debug, Clone, Error)]
@ -307,6 +312,69 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
}
Ok(())
}
/// Performs a transparent re-initialization handshake after a session-expired 404.
///
/// Takes an owned clone of the client (avoiding `&self` across `.await` so the
/// future remains `Send` without requiring `C: Sync`). POSTs the saved
/// initialize request without a session ID, extracts the new session ID and
/// protocol version, sends `notifications/initialized`, and returns the new
/// `(session_id, protocol_headers)` pair. The init result message is **not**
/// forwarded to the handler because the handler already processed the original
/// initialization.
async fn perform_reinitialization(
client: C,
saved_init_request: ClientJsonRpcMessage,
uri: Arc<str>,
auth_header: Option<String>,
custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<(Option<Arc<str>>, HashMap<HeaderName, HeaderValue>), StreamableHttpError<C::Error>>
{
let (init_msg, new_session_id_str) = client
.post_message(
uri.clone(),
saved_init_request,
None,
auth_header.clone(),
custom_headers.clone(),
)
.await?
.expect_initialized::<C::Error>()
.await?;
let new_session_id: Option<Arc<str>> = new_session_id_str.map(|s| Arc::from(s.as_str()));
// Start from custom_headers, then inject the negotiated MCP-Protocol-Version
// so all subsequent requests carry the right version (MCP 2025-06-18 spec).
let mut new_protocol_headers = custom_headers;
if let ServerJsonRpcMessage::Response(response) = &init_msg {
if let ServerResult::InitializeResult(init_result) = &response.result {
if let Ok(hv) = HeaderValue::from_str(init_result.protocol_version.as_str()) {
new_protocol_headers
.insert(HeaderName::from_static("mcp-protocol-version"), hv);
}
}
}
let initialized_notification = ClientJsonRpcMessage::notification(
ClientNotification::InitializedNotification(InitializedNotification {
method: Default::default(),
extensions: Default::default(),
}),
);
client
.post_message(
uri,
initialized_notification,
new_session_id.clone(),
auth_header,
new_protocol_headers.clone(),
)
.await?
.expect_accepted_or_json::<C::Error>()?;
Ok((new_session_id, new_protocol_headers))
}
}
impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
@ -338,14 +406,15 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
responder,
message: initialize_request,
} = context.recv_from_handler().await?;
let saved_init_request = initialize_request.clone();
let (message, session_id) = match self
.client
.post_message(
config.uri.clone(),
initialize_request,
None,
self.config.auth_header,
self.config.custom_headers,
config.auth_header.clone(),
config.custom_headers.clone(),
)
.await
{
@ -364,7 +433,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
));
}
};
let session_id: Option<Arc<str>> = if let Some(session_id) = session_id {
let mut session_id: Option<Arc<str>> = if let Some(session_id) = session_id {
Some(session_id.into())
} else {
if !self.config.allow_stateless {
@ -378,7 +447,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
// Extract the negotiated protocol version from the init response
// and build a custom headers map that includes MCP-Protocol-Version
// for all subsequent HTTP requests (per MCP 2025-06-18 spec).
let protocol_headers = {
let mut protocol_headers = {
let mut headers = config.custom_headers.clone();
if let ServerJsonRpcMessage::Response(response) = &message {
if let ServerResult::InitializeResult(init_result) = &response.result {
@ -392,7 +461,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
};
// Store session info for cleanup when run() exits (not spawned, so cleanup completes before close() returns)
let session_cleanup_info = session_id.as_ref().map(|sid| SessionCleanupInfo {
let mut session_cleanup_info = session_id.as_ref().map(|sid| SessionCleanupInfo {
client: self.client.clone(),
uri: config.uri.clone(),
session_id: sid.clone(),
@ -516,17 +585,171 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
match event {
Event::ClientMessage(send_request) => {
let WorkerSendRequest { message, responder } = send_request;
// Pass a clone to the first attempt so `message` is retained for a
// potential re-init retry. `post_message` takes ownership and the
// trait cannot be changed, so the clone is unavoidable.
let response = self
.client
.post_message(
config.uri.clone(),
message,
message.clone(),
session_id.clone(),
config.auth_header.clone(),
protocol_headers.clone(),
)
.await;
let send_result = match response {
Err(StreamableHttpError::SessionExpired) => {
// The server discarded the session (HTTP 404). Perform a
// fresh handshake once and replay the original message.
tracing::info!(
"session expired (HTTP 404), attempting transparent re-initialization"
);
match Self::perform_reinitialization(
self.client.clone(),
saved_init_request.clone(),
config.uri.clone(),
config.auth_header.clone(),
config.custom_headers.clone(),
)
.await
{
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;
protocol_headers = new_protocol_headers;
session_cleanup_info =
session_id.as_ref().map(|sid| SessionCleanupInfo {
client: self.client.clone(),
uri: config.uri.clone(),
session_id: sid.clone(),
auth_header: config.auth_header.clone(),
protocol_headers: protocol_headers.clone(),
});
if let Some(new_sid) = &session_id {
let client = self.client.clone();
let uri = config.uri.clone();
let new_sid = new_sid.clone();
let auth_header = config.auth_header.clone();
let retry_config = self.config.retry_config.clone();
let sse_tx = sse_worker_tx.clone();
let task_ct = transport_task_ct.clone();
let config_uri = config.uri.clone();
let config_auth = config.auth_header.clone();
let spawn_headers = protocol_headers.clone();
streams.spawn(async move {
match client
.get_stream(
uri,
new_sid.clone(),
None,
auth_header.clone(),
spawn_headers.clone(),
)
.await
{
Ok(stream) => {
let sse_stream = SseAutoReconnectStream::new(
stream,
StreamableHttpClientReconnect {
client: client.clone(),
session_id: new_sid,
uri: config_uri,
auth_header: config_auth,
custom_headers: spawn_headers,
},
retry_config,
);
Self::execute_sse_stream(
sse_stream,
sse_tx,
false,
task_ct.child_token(),
)
.await
}
Err(StreamableHttpError::ServerDoesNotSupportSse) => {
tracing::debug!(
"server doesn't support sse after re-init"
);
Ok(())
}
Err(e) => {
tracing::error!(
"fail to get common stream after re-init: {e}"
);
Err(e)
}
}
});
}
let retry_response = self
.client
.post_message(
config.uri.clone(),
message,
session_id.clone(),
config.auth_header.clone(),
protocol_headers.clone(),
)
.await;
match retry_response {
Err(e) => Err(e),
Ok(StreamableHttpPostResponse::Accepted) => {
tracing::trace!(
"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(
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(),
));
}
tracing::trace!("got new sse stream after re-init");
Ok(())
}
}
}
Err(reinit_err) => Err(reinit_err),
}
}
Err(e) => Err(e),
Ok(StreamableHttpPostResponse::Accepted) => {
tracing::trace!("client message accepted");

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

@ -7,9 +7,13 @@
use std::{collections::HashMap, sync::Arc};
use rmcp::{
ServiceExt,
model::{ClientJsonRpcMessage, ClientRequest, PingRequest, RequestId},
transport::{
streamable_http_client::{StreamableHttpClient, StreamableHttpError},
StreamableHttpClientTransport,
streamable_http_client::{
StreamableHttpClient, StreamableHttpClientTransportConfig, StreamableHttpError,
},
streamable_http_server::{
StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
},
@ -76,18 +80,10 @@ async fn test_stale_session_id_returns_status_aware_error() -> anyhow::Result<()
assert_eq!(raw_response.status(), reqwest::StatusCode::NOT_FOUND);
match result {
Err(StreamableHttpError::UnexpectedServerResponse(message)) => {
let message = message.to_string();
assert!(
message.contains("404"),
"error should include HTTP status code, got: {message}"
);
assert!(
message.to_ascii_lowercase().contains("session not found"),
"error should include session-not-found hint, got: {message}"
);
Err(StreamableHttpError::SessionExpired) => {
// Expected: post_message detects 404 with a session ID and returns SessionExpired
}
other => panic!("expected UnexpectedServerResponse, got: {other:?}"),
other => panic!("expected SessionExpired, got: {other:?}"),
}
ct.cancel();
@ -95,3 +91,82 @@ async fn test_stale_session_id_returns_status_aware_error() -> anyhow::Result<()
Ok(())
}
/// Verify that when the server loses a session (returns HTTP 404), the client
/// transparently re-initializes and the original request succeeds.
#[tokio::test]
async fn test_transparent_reinitialization_on_session_expiry() -> 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;
}
});
// Connect a full client transport (this performs initialize + notifications/initialized)
let transport = StreamableHttpClientTransport::from_config(
StreamableHttpClientTransportConfig::with_uri(format!("http://{addr}/mcp")),
);
let client = ().serve(transport).await?;
// Verify the session is established: list_all_resources() succeeds
let _resources = client.list_all_resources().await?;
// Capture the current session ID from the server
let original_session_id = {
let sessions = session_manager.sessions.read().await;
sessions
.keys()
.next()
.cloned()
.expect("session should exist")
};
// 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 trigger transparent re-initialization and still succeed
let _resources_after = client.list_all_resources().await?;
// Verify the server created a new session with a different ID
{
let sessions = session_manager.sessions.read().await;
let new_session_id = sessions
.keys()
.next()
.expect("new session should exist after re-initialization");
assert_ne!(
new_session_id, &original_session_id,
"new session ID should differ from the original"
);
}
let _ = client.cancel().await;
ct.cancel();
server_handle.await?;
Ok(())
}