From ca4564a46a90d626a397156b370488aecf1557fa Mon Sep 17 00:00:00 2001 From: Dale Seo <5466341+DaleSeo@users.noreply.github.com> Date: Sun, 6 Sep 2026 18:11:53 +0900 Subject: [PATCH] feat: add ServerHandler::negotiate_initialize --- crates/rmcp/src/handler/server.rs | 66 ++++++++++-- crates/rmcp/src/model.rs | 84 ++++++++++++++- crates/rmcp/src/service/server.rs | 4 +- .../test_protocol_version_negotiation.rs | 100 +++++++++++++++++- examples/servers/src/common/counter.rs | 5 +- 5 files changed, 246 insertions(+), 13 deletions(-) diff --git a/crates/rmcp/src/handler/server.rs b/crates/rmcp/src/handler/server.rs index b5f70be3f..70985527c 100644 --- a/crates/rmcp/src/handler/server.rs +++ b/crates/rmcp/src/handler/server.rs @@ -321,16 +321,59 @@ macro_rules! server_handler_methods { context: RequestContext, ) -> impl Future> + MaybeSendFuture + '_ { context.peer.set_peer_info(request.clone()); + std::future::ready(self.negotiate_initialize(&request)) + } + /// Build the `initialize` response for `request`, negotiating the + /// protocol version against [`Self::supported_protocol_versions`]. + /// + /// This is the whole body of the default [`Self::initialize`] minus its + /// `set_peer_info` side effect, so a server that overrides `initialize` + /// to add its own can call this instead of restating the negotiation + /// rule: + /// + /// ``` + /// use rmcp::{ + /// ErrorData as McpError, RoleServer, ServerHandler, + /// model::{InitializeRequestParams, InitializeResult, ServerInfo}, + /// service::RequestContext, + /// }; + /// + /// struct MyServer; + /// + /// impl ServerHandler for MyServer { + /// fn get_info(&self) -> ServerInfo { + /// ServerInfo::default() + /// } + /// + /// async fn initialize( + /// &self, + /// request: InitializeRequestParams, + /// context: RequestContext, + /// ) -> Result { + /// // ... record telemetry, register the peer, etc. + /// context.peer.set_peer_info(request.clone()); + /// self.negotiate_initialize(&request) + /// } + /// } + /// ``` + /// + /// # Errors + /// + /// Returns [`ErrorCode::UNSUPPORTED_PROTOCOL_VERSION`] when this server + /// supports no version that still has an `initialize` handshake. + /// + /// [`ErrorCode::UNSUPPORTED_PROTOCOL_VERSION`]: crate::model::ErrorCode::UNSUPPORTED_PROTOCOL_VERSION + fn negotiate_initialize( + &self, + request: &InitializeRequestParams, + ) -> Result { let mut info = self.get_info(); - let negotiated = negotiate_protocol_version( + info.protocol_version = negotiate_protocol_version( &request.protocol_version, std::mem::take(&mut info.protocol_version), &self.supported_protocol_versions(), - ); - std::future::ready(negotiated.map(|version| { - info.protocol_version = version; - info - })) + )?; + Ok(info) } /// Return the protocol versions supported by this server. /// @@ -339,6 +382,10 @@ macro_rules! server_handler_methods { /// list is advertised by [`Self::discover`], bounds what `initialize` /// negotiation may agree to, and is what per-request versions are /// validated against. + /// + /// To support everything up to some ceiling, use + /// [`ProtocolVersion::known_up_to`] rather than filtering + /// [`ProtocolVersion::KNOWN_VERSIONS`] by hand. fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { Cow::Borrowed(ProtocolVersion::KNOWN_VERSIONS) } @@ -621,6 +668,13 @@ macro_rules! impl_server_handler_for_wrapper { (**self).initialize(request, context) } + fn negotiate_initialize( + &self, + request: &InitializeRequestParams, + ) -> Result { + (**self).negotiate_initialize(request) + } + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { (**self).supported_protocol_versions() } diff --git a/crates/rmcp/src/model.rs b/crates/rmcp/src/model.rs index 6a1409870..be911f7f9 100644 --- a/crates/rmcp/src/model.rs +++ b/crates/rmcp/src/model.rs @@ -177,7 +177,7 @@ impl ProtocolVersion { /// First protocol version that requires SEP-2243 standard HTTP headers. pub const STANDARD_HEADERS: Self = Self::V_2026_07_28; - /// All protocol versions known to this SDK. + /// All protocol versions known to this SDK, oldest first. pub const KNOWN_VERSIONS: &[Self] = &[ Self::V_2024_11_05, Self::V_2025_03_26, @@ -190,6 +190,43 @@ impl ProtocolVersion { pub fn as_str(&self) -> &str { &self.0 } + + /// The known versions up to and including `max`, oldest first. + /// + /// Servers that implement every revision up to some ceiling can return + /// this from `supported_protocol_versions` instead of filtering + /// [`Self::KNOWN_VERSIONS`] by hand. `max` itself need not be a known + /// version; the result is empty when it predates all of them. + /// + /// The result borrows from [`Self::KNOWN_VERSIONS`], so call it directly + /// in the method body — it needs no `static` and no `LazyLock`: + /// + /// ```rust,ignore + /// const MAX_SUPPORTED: ProtocolVersion = ProtocolVersion::V_2025_11_25; + /// + /// fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + /// Cow::Borrowed(ProtocolVersion::known_up_to(&MAX_SUPPORTED)) + /// } + /// ``` + /// + /// ``` + /// # use rmcp::model::ProtocolVersion; + /// assert_eq!( + /// ProtocolVersion::known_up_to(&ProtocolVersion::V_2025_06_18), + /// &[ + /// ProtocolVersion::V_2024_11_05, + /// ProtocolVersion::V_2025_03_26, + /// ProtocolVersion::V_2025_06_18, + /// ], + /// ); + /// ``` + pub fn known_up_to(max: &Self) -> &'static [Self] { + let count = Self::KNOWN_VERSIONS + .iter() + .take_while(|version| version.as_str() <= max.as_str()) + .count(); + &Self::KNOWN_VERSIONS[..count] + } } impl Serialize for ProtocolVersion { @@ -4643,6 +4680,51 @@ mod tests { use super::*; + #[test] + fn known_versions_are_ordered_oldest_first() { + // `known_up_to` walks the list as a sorted prefix. + assert!( + ProtocolVersion::KNOWN_VERSIONS + .windows(2) + .all(|pair| pair[0].as_str() < pair[1].as_str()) + ); + } + + #[test] + fn known_up_to_includes_the_ceiling_itself() { + assert_eq!( + ProtocolVersion::known_up_to(&ProtocolVersion::V_2024_11_05), + &[ProtocolVersion::V_2024_11_05] + ); + } + + #[test] + fn known_up_to_the_newest_version_yields_every_known_version() { + assert_eq!( + ProtocolVersion::known_up_to(&ProtocolVersion::V_2026_07_28), + ProtocolVersion::KNOWN_VERSIONS + ); + } + + #[test] + fn known_up_to_an_unknown_ceiling_stops_at_the_versions_below_it() { + let unknown = ProtocolVersion(Cow::Borrowed("2025-07-01")); + assert_eq!( + ProtocolVersion::known_up_to(&unknown), + &[ + ProtocolVersion::V_2024_11_05, + ProtocolVersion::V_2025_03_26, + ProtocolVersion::V_2025_06_18, + ] + ); + } + + #[test] + fn known_up_to_a_ceiling_below_every_known_version_is_empty() { + let ancient = ProtocolVersion(Cow::Borrowed("1999-01-01")); + assert!(ProtocolVersion::known_up_to(&ancient).is_empty()); + } + #[cfg(feature = "transport-streamable-http-client")] #[test] fn transport_closed_marker_accepts_only_the_process_local_token() { diff --git a/crates/rmcp/src/service/server.rs b/crates/rmcp/src/service/server.rs index 70a148641..29e46907a 100644 --- a/crates/rmcp/src/service/server.rs +++ b/crates/rmcp/src/service/server.rs @@ -499,7 +499,9 @@ pub(crate) fn negotiate_protocol_version( server_supported, )); }; - tracing::warn!( + // Falling back is the designed answer for a pinned client, and stateless + // HTTP re-runs it on every request, so this is not a warning. + tracing::debug!( client_requested = %client_requested, server_fallback = %legacy_fallback, "client requested a protocol version unavailable over initialize; falling back to server default" diff --git a/crates/rmcp/tests/test_protocol_version_negotiation.rs b/crates/rmcp/tests/test_protocol_version_negotiation.rs index 717a63f18..aa4607a1c 100644 --- a/crates/rmcp/tests/test_protocol_version_negotiation.rs +++ b/crates/rmcp/tests/test_protocol_version_negotiation.rs @@ -5,13 +5,19 @@ #![cfg(not(feature = "local"))] #![cfg(feature = "client")] -use std::borrow::Cow; +use std::{ + borrow::Cow, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, +}; use rmcp::{ ClientHandler, ErrorData, RoleServer, ServerHandler, ServiceExt, model::{ - ClientInfo, ErrorCode, InitializeRequestParams, InitializeResult, ProtocolVersion, - ServerInfo, + ClientCapabilities, ClientInfo, ErrorCode, Implementation, InitializeRequestParams, + InitializeResult, ProtocolVersion, ServerInfo, }, service::{ClientInitializeError, RequestContext}, }; @@ -232,3 +238,91 @@ async fn narrowed_server_caps_even_when_it_overrides_initialize() { "the handshake layer should not raise the version above what the server supports" ); } + +/// Overrides `initialize` to run a side effect, then delegates the version +/// answer back to the SDK with [`ServerHandler::negotiate_initialize`]. +#[derive(Debug, Clone, Default)] +struct DelegatingServer { + initializations: Arc, +} + +impl ServerHandler for DelegatingServer { + fn get_info(&self) -> ServerInfo { + ServerInfo::default() + } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Borrowed(HANDSHAKE_VERSIONS) + } + + async fn initialize( + &self, + request: InitializeRequestParams, + context: RequestContext, + ) -> Result { + self.initializations.fetch_add(1, Ordering::Relaxed); + context.peer.set_peer_info(request.clone()); + self.negotiate_initialize(&request) + } +} + +fn initialize_params(protocol_version: ProtocolVersion) -> InitializeRequestParams { + let mut params = InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("test-client", "0.0.0"), + ); + params.protocol_version = protocol_version; + params +} + +#[test] +fn negotiate_initialize_echoes_a_supported_version() { + let result = NarrowedServer + .negotiate_initialize(&initialize_params(ProtocolVersion::V_2025_06_18)) + .expect("a supported handshake version should negotiate"); + assert_eq!(result.protocol_version, ProtocolVersion::V_2025_06_18); +} + +#[test] +fn negotiate_initialize_caps_at_supported_versions() { + let result = NarrowedServer + .negotiate_initialize(&initialize_params(ProtocolVersion::V_2026_07_28)) + .expect("an unsupported version should fall back rather than fail"); + assert_eq!(result.protocol_version, ProtocolVersion::V_2025_11_25); +} + +#[test] +fn negotiate_initialize_keeps_the_rest_of_get_info() { + let server = NarrowedServer; + let result = server + .negotiate_initialize(&initialize_params(ProtocolVersion::V_2026_07_28)) + .expect("an unsupported version should fall back rather than fail"); + assert_eq!(result.capabilities, server.get_info().capabilities); +} + +#[test] +fn negotiate_initialize_rejects_when_no_handshake_version_is_supported() { + let error = ModernOnlyServer + .negotiate_initialize(&initialize_params(ProtocolVersion::V_2026_07_28)) + .expect_err("a server with no handshake version cannot answer initialize"); + assert_eq!(error.code, ErrorCode::UNSUPPORTED_PROTOCOL_VERSION); +} + +#[tokio::test] +async fn delegating_server_negotiates_like_the_default_initialize() { + let negotiated = + negotiated_version_with(DelegatingServer::default(), ProtocolVersion::V_2026_07_28).await; + assert_eq!( + negotiated, + ProtocolVersion::V_2025_11_25, + "an override that delegates should answer what the default initialize would" + ); +} + +#[tokio::test] +async fn delegating_server_still_runs_its_own_side_effect() { + let server = DelegatingServer::default(); + let initializations = Arc::clone(&server.initializations); + negotiated_version_with(server, ProtocolVersion::V_2025_06_18).await; + assert_eq!(initializations.load(Ordering::Relaxed), 1); +} diff --git a/examples/servers/src/common/counter.rs b/examples/servers/src/common/counter.rs index c6602770f..3bd16173c 100644 --- a/examples/servers/src/common/counter.rs +++ b/examples/servers/src/common/counter.rs @@ -270,7 +270,7 @@ impl ServerHandler for Counter { async fn initialize( &self, - _request: InitializeRequestParams, + request: InitializeRequestParams, context: RequestContext, ) -> Result { if let Some(http_request_part) = context.extensions.get::() { @@ -278,7 +278,8 @@ impl ServerHandler for Counter { let initialize_uri = &http_request_part.uri; tracing::info!(?initialize_headers, %initialize_uri, "initialize from http server"); } - Ok(self.get_info()) + context.peer.set_peer_info(request.clone()); + self.negotiate_initialize(&request) } }