feat: add ServerHandler::negotiate_initialize (#1247)
This commit is contained in:
parent
7d14f3a8a2
commit
f9a68bfa5c
5 changed files with 246 additions and 13 deletions
|
|
@ -321,16 +321,59 @@ macro_rules! server_handler_methods {
|
||||||
context: RequestContext<RoleServer>,
|
context: RequestContext<RoleServer>,
|
||||||
) -> impl Future<Output = Result<InitializeResult, McpError>> + MaybeSendFuture + '_ {
|
) -> impl Future<Output = Result<InitializeResult, McpError>> + MaybeSendFuture + '_ {
|
||||||
context.peer.set_peer_info(request.clone());
|
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<RoleServer>,
|
||||||
|
/// ) -> Result<InitializeResult, McpError> {
|
||||||
|
/// // ... 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<InitializeResult, McpError> {
|
||||||
let mut info = self.get_info();
|
let mut info = self.get_info();
|
||||||
let negotiated = negotiate_protocol_version(
|
info.protocol_version = negotiate_protocol_version(
|
||||||
&request.protocol_version,
|
&request.protocol_version,
|
||||||
std::mem::take(&mut info.protocol_version),
|
std::mem::take(&mut info.protocol_version),
|
||||||
&self.supported_protocol_versions(),
|
&self.supported_protocol_versions(),
|
||||||
);
|
)?;
|
||||||
std::future::ready(negotiated.map(|version| {
|
Ok(info)
|
||||||
info.protocol_version = version;
|
|
||||||
info
|
|
||||||
}))
|
|
||||||
}
|
}
|
||||||
/// Return the protocol versions supported by this server.
|
/// 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`
|
/// list is advertised by [`Self::discover`], bounds what `initialize`
|
||||||
/// negotiation may agree to, and is what per-request versions are
|
/// negotiation may agree to, and is what per-request versions are
|
||||||
/// validated against.
|
/// 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]> {
|
fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> {
|
||||||
Cow::Borrowed(ProtocolVersion::KNOWN_VERSIONS)
|
Cow::Borrowed(ProtocolVersion::KNOWN_VERSIONS)
|
||||||
}
|
}
|
||||||
|
|
@ -621,6 +668,13 @@ macro_rules! impl_server_handler_for_wrapper {
|
||||||
(**self).initialize(request, context)
|
(**self).initialize(request, context)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn negotiate_initialize(
|
||||||
|
&self,
|
||||||
|
request: &InitializeRequestParams,
|
||||||
|
) -> Result<InitializeResult, McpError> {
|
||||||
|
(**self).negotiate_initialize(request)
|
||||||
|
}
|
||||||
|
|
||||||
fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> {
|
fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> {
|
||||||
(**self).supported_protocol_versions()
|
(**self).supported_protocol_versions()
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -177,7 +177,7 @@ impl ProtocolVersion {
|
||||||
/// First protocol version that requires SEP-2243 standard HTTP headers.
|
/// First protocol version that requires SEP-2243 standard HTTP headers.
|
||||||
pub const STANDARD_HEADERS: Self = Self::V_2026_07_28;
|
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] = &[
|
pub const KNOWN_VERSIONS: &[Self] = &[
|
||||||
Self::V_2024_11_05,
|
Self::V_2024_11_05,
|
||||||
Self::V_2025_03_26,
|
Self::V_2025_03_26,
|
||||||
|
|
@ -190,6 +190,43 @@ impl ProtocolVersion {
|
||||||
pub fn as_str(&self) -> &str {
|
pub fn as_str(&self) -> &str {
|
||||||
&self.0
|
&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 {
|
impl Serialize for ProtocolVersion {
|
||||||
|
|
@ -4643,6 +4680,51 @@ mod tests {
|
||||||
|
|
||||||
use super::*;
|
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")]
|
#[cfg(feature = "transport-streamable-http-client")]
|
||||||
#[test]
|
#[test]
|
||||||
fn transport_closed_marker_accepts_only_the_process_local_token() {
|
fn transport_closed_marker_accepts_only_the_process_local_token() {
|
||||||
|
|
|
||||||
|
|
@ -499,7 +499,9 @@ pub(crate) fn negotiate_protocol_version(
|
||||||
server_supported,
|
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,
|
client_requested = %client_requested,
|
||||||
server_fallback = %legacy_fallback,
|
server_fallback = %legacy_fallback,
|
||||||
"client requested a protocol version unavailable over initialize; falling back to server default"
|
"client requested a protocol version unavailable over initialize; falling back to server default"
|
||||||
|
|
|
||||||
|
|
@ -5,13 +5,19 @@
|
||||||
#![cfg(not(feature = "local"))]
|
#![cfg(not(feature = "local"))]
|
||||||
#![cfg(feature = "client")]
|
#![cfg(feature = "client")]
|
||||||
|
|
||||||
use std::borrow::Cow;
|
use std::{
|
||||||
|
borrow::Cow,
|
||||||
|
sync::{
|
||||||
|
Arc,
|
||||||
|
atomic::{AtomicUsize, Ordering},
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
use rmcp::{
|
use rmcp::{
|
||||||
ClientHandler, ErrorData, RoleServer, ServerHandler, ServiceExt,
|
ClientHandler, ErrorData, RoleServer, ServerHandler, ServiceExt,
|
||||||
model::{
|
model::{
|
||||||
ClientInfo, ErrorCode, InitializeRequestParams, InitializeResult, ProtocolVersion,
|
ClientCapabilities, ClientInfo, ErrorCode, Implementation, InitializeRequestParams,
|
||||||
ServerInfo,
|
InitializeResult, ProtocolVersion, ServerInfo,
|
||||||
},
|
},
|
||||||
service::{ClientInitializeError, RequestContext},
|
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"
|
"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<AtomicUsize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
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<RoleServer>,
|
||||||
|
) -> Result<InitializeResult, ErrorData> {
|
||||||
|
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);
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -270,7 +270,7 @@ impl ServerHandler for Counter {
|
||||||
|
|
||||||
async fn initialize(
|
async fn initialize(
|
||||||
&self,
|
&self,
|
||||||
_request: InitializeRequestParams,
|
request: InitializeRequestParams,
|
||||||
context: RequestContext<RoleServer>,
|
context: RequestContext<RoleServer>,
|
||||||
) -> Result<InitializeResult, McpError> {
|
) -> Result<InitializeResult, McpError> {
|
||||||
if let Some(http_request_part) = context.extensions.get::<axum::http::request::Parts>() {
|
if let Some(http_request_part) = context.extensions.get::<axum::http::request::Parts>() {
|
||||||
|
|
@ -278,7 +278,8 @@ impl ServerHandler for Counter {
|
||||||
let initialize_uri = &http_request_part.uri;
|
let initialize_uri = &http_request_part.uri;
|
||||||
tracing::info!(?initialize_headers, %initialize_uri, "initialize from http server");
|
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)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue