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: negotiate stateless initialize versions (#1080)

Browse source
This commit is contained in:
Dale Seo 2026-07-29 15:47:19 -04:00 • committed by GitHub
commit 596a7e1e0d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 124 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -28,10 +28,13 @@ use crate::{
ClientCapabilities, ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorCode,
ErrorData, GetExtensions, GetMeta, Implementation, InitializeRequest,
InitializeRequestParams, InitializedNotification, JsonObject, JsonRpcError,
ProtocolVersion, RequestId, ServerJsonRpcMessage,
ProtocolVersion, RequestId, ServerInfo, ServerJsonRpcMessage, ServerResult,
},
serve_server,
service::{serve_directly_with_ct, uses_legacy_lifecycle},
service::{
NotificationContext, RequestContext, Service, negotiate_protocol_version,
serve_directly_with_ct, uses_legacy_lifecycle,
},
transport::{
OneshotTransport, TransportAdapterIdentity,
common::{
@ -257,6 +260,49 @@ fn message_has_per_request_protocol_version(message: &ClientJsonRpcMessage) -> b
}
}
struct NegotiatingStatelessHttpService<S>(S);
impl<S: Service<RoleServer>> Service<RoleServer> for NegotiatingStatelessHttpService<S> {
async fn handle_request(
&self,
request: ClientRequest,
context: RequestContext<RoleServer>,
) -> Result<ServerResult, ErrorData> {
let requested_protocol_version =
if let ClientRequest::InitializeRequest(initialize) = &request {
Some(initialize.params.protocol_version.clone())
} else {
None
};
let peer = context.peer.clone();
let mut response = self.0.handle_request(request, context).await?;
if let (Some(requested), ServerResult::InitializeResult(result)) =
(requested_protocol_version, &mut response)
{
result.protocol_version =
negotiate_protocol_version(&requested, result.protocol_version.clone());
if let Some(peer_info) = peer.peer_info() {
let mut peer_info = (*peer_info).clone();
peer_info.protocol_version = result.protocol_version.clone();
peer.set_peer_info(peer_info);
}
}
Ok(response)
}
async fn handle_notification(
&self,
notification: ClientNotification,
context: NotificationContext<RoleServer>,
) -> Result<(), ErrorData> {
self.0.handle_notification(notification, context).await
}
fn get_info(&self) -> ServerInfo {
self.0.get_info()
}
}
#[expect(
clippy::result_large_err,
reason = "BoxResponse is intentionally large; matches other handlers in this file"
@ -1044,7 +1090,12 @@ where
// disconnect can cancel the in-flight handler (#857), as in the
// non-negotiated stateless path below.
let request_ct = CancellationToken::new();
let service = serve_directly_with_ct(service, transport, peer_info, request_ct.clone());
let service = serve_directly_with_ct(
NegotiatingStatelessHttpService(service),
transport,
peer_info,
request_ct.clone(),
);
tokio::spawn(async move {
let _ = service.waiting().await;
});
@ -1789,8 +1840,12 @@ where
// unpersisted response can cancel the in-flight handler on
// disconnect (#857).
let request_ct = CancellationToken::new();
let service =
serve_directly_with_ct(service, transport, peer_info, request_ct.clone());
let service = serve_directly_with_ct(
NegotiatingStatelessHttpService(service),
transport,
peer_info,
request_ct.clone(),
);
tokio::spawn(async move {
// on service created
let _ = service.waiting().await;

75
crates/rmcp/tests/test_stateless_protocol_version.rs
View file

@ -4,30 +4,51 @@
#![cfg(not(feature = "local"))]
use rmcp::{
model::ProtocolVersion,
ErrorData, RoleServer, ServerHandler,
model::{
InitializeRequestParams, InitializeResult, ProtocolVersion, ServerCapabilities, ServerInfo,
},
service::RequestContext,
transport::streamable_http_server::{
StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
},
};
use tokio_util::sync::CancellationToken;
mod common;
use common::calculator::Calculator;
#[derive(Clone)]
struct OverridingInitialize;
fn stateless_json_config() -> StreamableHttpServerConfig {
impl ServerHandler for OverridingInitialize {
fn get_info(&self) -> ServerInfo {
ServerInfo::new(ServerCapabilities::default())
}
async fn initialize(
&self,
_request: InitializeRequestParams,
_context: RequestContext<RoleServer>,
) -> Result<InitializeResult, ErrorData> {
Ok(self.get_info())
}
}
fn stateless_sse_config() -> StreamableHttpServerConfig {
StreamableHttpServerConfig::default()
.with_legacy_session_mode(false)
.with_json_response(true)
.with_sse_keep_alive(None)
.with_cancellation_token(CancellationToken::new())
}
fn stateless_json_config() -> StreamableHttpServerConfig {
stateless_sse_config().with_json_response(true)
}
async fn spawn_server(
config: StreamableHttpServerConfig,
) -> (reqwest::Client, String, CancellationToken) {
let ct = config.cancellation_token.clone();
let service: StreamableHttpService<Calculator, LocalSessionManager> =
StreamableHttpService::new(|| Ok(Calculator::new()), Default::default(), config);
let service: StreamableHttpService<OverridingInitialize, LocalSessionManager> =
StreamableHttpService::new(|| Ok(OverridingInitialize), Default::default(), config);
let router = axum::Router::new().nest_service("/mcp", service);
let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@ -65,11 +86,27 @@ async fn post_init(client: &reqwest::Client, url: &str, body_version: &str) -> s
.await
.expect("send request");
assert!(resp.status().is_success(), "HTTP {}", resp.status());
resp.json().await.expect("parse JSON")
let is_json = resp
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.is_some_and(|value| value.starts_with("application/json"));
if is_json {
resp.json().await.expect("parse JSON")
} else {
let body = resp.text().await.expect("read SSE body");
let data = body
.lines()
.find_map(|line| line.strip_prefix("data:"))
.map(str::trim)
.filter(|data| !data.is_empty())
.expect("SSE response contains data");
serde_json::from_str(data).expect("parse SSE data")
}
}
#[tokio::test]
async fn stateless_init_echoes_known_version() {
async fn stateless_json_init_echoes_known_versions_when_handler_overrides_initialize() {
let (client, url, ct) = spawn_server(stateless_json_config()).await;
for version in ProtocolVersion::KNOWN_VERSIONS {
@ -85,14 +122,30 @@ async fn stateless_init_echoes_known_version() {
}
#[tokio::test]
async fn stateless_init_unknown_version_falls_back_to_latest() {
async fn stateless_sse_init_echoes_known_versions_when_handler_overrides_initialize() {
let (client, url, ct) = spawn_server(stateless_sse_config()).await;
for version in ProtocolVersion::KNOWN_VERSIONS {
let resp = post_init(&client, &url, version.as_str()).await;
assert_eq!(
resp["result"]["protocolVersion"],
version.as_str(),
"known version {version} should be echoed back"
);
}
ct.cancel();
}
#[tokio::test]
async fn stateless_json_init_preserves_handler_fallback_for_unknown_version() {
let (client, url, ct) = spawn_server(stateless_json_config()).await;
let resp = post_init(&client, &url, "1999-01-01").await;
assert_eq!(
resp["result"]["protocolVersion"],
ProtocolVersion::LATEST.as_str(),
"unknown version should fall back to LATEST"
"unknown version should preserve the handler's fallback"
);
ct.cancel();