fix: negotiate stateless initialize versions (#1080)

crates/rmcp/src/transport/streamable_http_server/tower.rs

@@ -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 =
+                    let service = serve_directly_with_ct(
-                        serve_directly_with_ct(service, transport, peer_info, request_ct.clone());
+                        NegotiatingStatelessHttpService(service),
+                        transport,
+                        peer_info,
+                        request_ct.clone(),
+                    );
                     tokio::spawn(async move {
                         // on service created
                         let _ = service.waiting().await;

crates/rmcp/tests/test_stateless_protocol_version.rs

@@ -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;
+#[derive(Clone)]
-use common::calculator::Calculator;
+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> =
+    let service: StreamableHttpService<OverridingInitialize, LocalSessionManager> =
-        StreamableHttpService::new(|| Ok(Calculator::new()), Default::default(), config);
+        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();