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
1 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;