diff --git a/crates/rmcp/src/service.rs b/crates/rmcp/src/service.rs index aaec82895..34e6e92ba 100644 --- a/crates/rmcp/src/service.rs +++ b/crates/rmcp/src/service.rs @@ -161,7 +161,7 @@ pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone { /// Rejects outbound requests when the negotiated protocol forbids this role /// from sending any. #[doc(hidden)] - fn enforce_outbound_request(_peer_info: Option<&Self::PeerInfo>) -> Result<(), ServiceError> { + fn enforce_outbound_request(_peer: &Peer) -> Result<(), ServiceError> { Ok(()) } @@ -226,6 +226,18 @@ pub(crate) fn peer_request_association( tokio::task_local! { pub(crate) static ORIGINATING_REQUEST: RequestId; + /// The protocol version the originating request names in its `_meta`. + pub(crate) static ORIGINATING_PROTOCOL_VERSION: Option; +} + +/// The protocol version named by the request the current handler is serving, +/// if any. Like [`ORIGINATING_REQUEST`], it does not cross `tokio::spawn`. +#[cfg(feature = "server")] +pub(crate) fn originating_protocol_version() -> Option { + ORIGINATING_PROTOCOL_VERSION + .try_with(Clone::clone) + .ok() + .flatten() } /// Marker in an outbound request's non-serialized [`Extensions`] identifying @@ -717,7 +729,7 @@ pub struct Peer { progress_timeout_watchers: ProgressTimeoutWatchers, info: Arc>>>, client_request_metadata: Arc>, - request_metadata_required: Arc, + inline_protocol_version: Arc>, subscription_channels: Arc>>, #[cfg(feature = "client")] response_cache: client::cache::PeerResponseCache, @@ -735,7 +747,7 @@ where progress_timeout_watchers: self.progress_timeout_watchers.clone(), info: self.info.clone(), client_request_metadata: self.client_request_metadata.clone(), - request_metadata_required: self.request_metadata_required.clone(), + inline_protocol_version: self.inline_protocol_version.clone(), subscription_channels: self.subscription_channels.clone(), #[cfg(feature = "client")] response_cache: self.response_cache.clone(), @@ -811,7 +823,7 @@ impl Peer { progress_timeout_watchers: Default::default(), info: Arc::new(std::sync::RwLock::new(peer_info.map(Arc::new))), client_request_metadata: Default::default(), - request_metadata_required: Default::default(), + inline_protocol_version: Default::default(), subscription_channels: Default::default(), #[cfg(feature = "client")] response_cache: Default::default(), @@ -860,7 +872,7 @@ impl Peer { options: PeerRequestOptions, subscription_sender: Option>, ) -> Result, ServiceError> { - R::enforce_outbound_request(self.peer_info().as_deref())?; + R::enforce_outbound_request(self)?; if let Ok(originating) = ORIGINATING_REQUEST.try_with(|id| id.clone()) { request .extensions_mut() @@ -1024,16 +1036,22 @@ impl Peer { debug_assert!(result.is_ok(), "client request metadata set more than once"); } + /// Marks an inline lifecycle session that opened on `protocol_version`. + #[cfg(feature = "server")] + pub(crate) fn require_request_metadata(&self, protocol_version: ProtocolVersion) { + let _ = self.inline_protocol_version.set(protocol_version); + } + + /// The version an inline lifecycle session opened on. Such a session never + /// stores peer info. #[cfg(feature = "server")] - pub(crate) fn require_request_metadata(&self) { - self.request_metadata_required - .store(true, std::sync::atomic::Ordering::Release); + pub(crate) fn inline_protocol_version(&self) -> Option<&ProtocolVersion> { + self.inline_protocol_version.get() } #[cfg(feature = "server")] pub(crate) fn request_metadata_required(&self) -> bool { - self.request_metadata_required - .load(std::sync::atomic::Ordering::Acquire) + self.inline_protocol_version.get().is_some() } pub fn is_transport_closed(&self) -> bool { @@ -1586,9 +1604,16 @@ where }; let current_span = tracing::Span::current(); let handler_id = id.clone(); + let handler_version = context.meta.protocol_version(); spawn_service_task(async move { let result = ORIGINATING_REQUEST - .scope(handler_id, service.handle_request(request, context)) + .scope( + handler_id, + ORIGINATING_PROTOCOL_VERSION.scope( + handler_version, + service.handle_request(request, context), + ), + ) .await; let response = match result { Ok(result) => { diff --git a/crates/rmcp/src/service/server.rs b/crates/rmcp/src/service/server.rs index f52b1159f..185fd543e 100644 --- a/crates/rmcp/src/service/server.rs +++ b/crates/rmcp/src/service/server.rs @@ -51,8 +51,15 @@ impl ServiceRole for RoleServer { } } - fn enforce_outbound_request(peer_info: Option<&Self::PeerInfo>) -> Result<(), ServiceError> { - if peer_info.is_some_and(|info| info.protocol_version >= ProtocolVersion::V_2026_07_28) { + fn enforce_outbound_request(peer: &Peer) -> Result<(), ServiceError> { + // Same order as `RequestContext::protocol_version`: the version the + // current request names wins. Outside a handler, fall back to the version + // an inline session (e.g. discover on stdio) opened on, since it never + // stores peer info, then to peer info. + let version = originating_protocol_version() + .or_else(|| peer.inline_protocol_version().cloned()) + .or_else(|| peer.peer_info().map(|info| info.protocol_version.clone())); + if version.is_some_and(|version| version >= ProtocolVersion::V_2026_07_28) { return Err(ServiceError::McpError(ErrorData::invalid_request( "server-to-client requests are not allowed on protocol 2026-07-28 or later; return InputRequiredResult instead", None, @@ -650,7 +657,7 @@ where })?; continue; } - peer.require_request_metadata(); + peer.require_request_metadata(requested_version); // Dispatch the request from inside the service loop rather than // inline: its handler may send notifications through `peer`, which // only complete once the loop drains `peer_rx`. @@ -1237,6 +1244,10 @@ impl Peer { where T: ElicitationSafe + for<'de> serde::Deserialize<'de>, { + // Reject on 2026-07-28 before the capability check, which would + // otherwise answer `CapabilityNotSupported` first. + RoleServer::enforce_outbound_request(self)?; + // Check if client supports form elicitation capability if !self .supported_elicitation_modes() @@ -1388,6 +1399,10 @@ impl Peer { elicitation_id: impl Into, timeout: Option, ) -> Result { + // Reject on 2026-07-28 before the capability check, which would + // otherwise answer `CapabilityNotSupported` first. + RoleServer::enforce_outbound_request(self)?; + // Check if client supports url elicitation if !self .supported_elicitation_modes() diff --git a/crates/rmcp/tests/test_sep_2260_request_association.rs b/crates/rmcp/tests/test_sep_2260_request_association.rs index 8c45296b9..16cf8c570 100644 --- a/crates/rmcp/tests/test_sep_2260_request_association.rs +++ b/crates/rmcp/tests/test_sep_2260_request_association.rs @@ -2,7 +2,8 @@ #![expect(deprecated, reason = "This test exercises the deprecated sampling API")] use rmcp::{ - ClientHandler, RoleClient, RoleServer, ServerHandler, ServiceError, ServiceExt, + ClientHandler, ClientLifecycleMode, ClientServiceExt, RoleClient, RoleServer, ServerHandler, + ServiceError, ServiceExt, model::{ CallToolRequestParams, CallToolResponse, CallToolResult, ClientConfig, ContentBlock, CreateMessageRequest, CreateMessageRequestParams, CreateMessageResult, ErrorCode, @@ -35,6 +36,14 @@ impl ServerHandler for SamplingServer { CreateMessageRequestParams::new(vec![SamplingMessage::user_text("nested")], 16); let outcome = match request.name.as_ref() { "sample" => context.peer.create_message(params).await.map(|_| ()), + // A spawned task loses the originating request, so the guard falls + // back to the session's version. + "sample_spawned" => { + let peer = context.peer.clone(); + tokio::spawn(async move { peer.create_message(params).await.map(|_| ()) }) + .await + .unwrap() + } "sample_generic" => context .peer .send_request(ServerRequest::CreateMessageRequest( @@ -47,6 +56,20 @@ impl ServerHandler for SamplingServer { .send_request(ServerRequest::PingRequest(PingRequest::default())) .await .map(|_| ()), + #[cfg(feature = "elicitation")] + "elicit_url" => match context + .peer + .elicit_url( + "visit", + url::Url::parse("https://example.com/").unwrap(), + "elicitation-1", + ) + .await + { + Ok(_) => Ok(()), + Err(rmcp::service::ElicitationError::Service(e)) => Err(e), + Err(e) => return Err(rmcp::ErrorData::internal_error(e.to_string(), None)), + }, other => panic!("unexpected tool {other}"), }; let text = match outcome { @@ -130,6 +153,46 @@ async fn call_tool_on_modern_pair(tool: &'static str) -> anyhow::Result Ok(text) } +/// Starts the pair the way stdio does: the server through plain `serve()`, the +/// client through the discover lifecycle on `version`. Unlike `serve_directly`, +/// this leaves the server's `peer_info` empty. +async fn call_tool_on_discover_pair( + tool: &'static str, + version: ProtocolVersion, +) -> anyhow::Result { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let server_handle = tokio::spawn(async move { + let server = SamplingServer.serve(server_transport).await?; + server.waiting().await?; + anyhow::Ok(()) + }); + let client = SamplingClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![version], + }, + ) + .await?; + + let result = client + .peer() + .call_tool(CallToolRequestParams::new(tool)) + .await?; + let text = result + .content + .first() + .unwrap() + .as_text() + .unwrap() + .text + .clone(); + + client.cancel().await?; + let _ = server_handle.await?; + Ok(text) +} + #[tokio::test] async fn sampling_from_handler_rejected_on_modern_protocol() -> anyhow::Result<()> { assert_eq!(call_tool_on_modern_pair("sample").await?, "rejected"); @@ -151,6 +214,139 @@ async fn ping_rejected_on_modern_protocol() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn sampling_from_handler_rejected_on_discover_lifecycle() -> anyhow::Result<()> { + assert_eq!( + call_tool_on_discover_pair("sample", ProtocolVersion::V_2026_07_28).await?, + "rejected" + ); + Ok(()) +} + +#[tokio::test] +async fn generic_send_request_rejected_on_discover_lifecycle() -> anyhow::Result<()> { + assert_eq!( + call_tool_on_discover_pair("sample_generic", ProtocolVersion::V_2026_07_28).await?, + "rejected" + ); + Ok(()) +} + +#[tokio::test] +async fn ping_rejected_on_discover_lifecycle() -> anyhow::Result<()> { + assert_eq!( + call_tool_on_discover_pair("ping", ProtocolVersion::V_2026_07_28).await?, + "rejected" + ); + Ok(()) +} + +#[tokio::test] +async fn spawned_sampling_rejected_on_discover_lifecycle() -> anyhow::Result<()> { + assert_eq!( + call_tool_on_discover_pair("sample_spawned", ProtocolVersion::V_2026_07_28).await?, + "rejected" + ); + Ok(()) +} + +/// Plays a raw client that skips discover and names a version in each request's +/// `_meta`, then calls the `sample` tool once per entry in `versions` on one +/// connection. Answers any `sampling/createMessage` the server sends. +async fn raw_inline_sample_calls(versions: &[&str]) -> anyhow::Result> { + let (server_io, client_io) = tokio::io::duplex(4096); + tokio::spawn(async move { + let server = SamplingServer.serve(server_io).await?; + server.waiting().await?; + anyhow::Ok(()) + }); + let (read, mut write) = tokio::io::split(client_io); + let mut lines = BufReader::new(read).lines(); + let mut texts = Vec::new(); + for (id, version) in versions.iter().enumerate() { + let call = json!({ + "jsonrpc": "2.0", + "id": id, + "method": "tools/call", + "params": { + "name": "sample", + "_meta": { + "io.modelcontextprotocol/protocolVersion": version, + "io.modelcontextprotocol/clientCapabilities": {} + } + } + }); + write.write_all(format!("{call}\n").as_bytes()).await?; + loop { + let line = tokio::time::timeout(std::time::Duration::from_secs(5), lines.next_line()) + .await?? + .expect("server message"); + let message: Value = serde_json::from_str(&line)?; + if message["method"] == "sampling/createMessage" { + let reply = json!({ + "jsonrpc": "2.0", + "id": message["id"], + "result": { + "role": "assistant", + "content": { "type": "text", "text": "pong" }, + "model": "test-model" + } + }); + write.write_all(format!("{reply}\n").as_bytes()).await?; + continue; + } + assert_eq!(message["id"], id, "unexpected message {message}"); + let text = message["result"]["content"][0]["text"] + .as_str() + .unwrap_or_else(|| panic!("tool result expected, got {message}")); + texts.push(text.to_owned()); + break; + } + } + Ok(texts) +} + +/// Each inline request names its own version, so the guard follows the request +/// a handler is serving, not the one that opened the session. +#[tokio::test] +async fn guard_follows_each_request_version_on_inline_session() -> anyhow::Result<()> { + assert_eq!( + raw_inline_sample_calls(&["2026-07-28", "2025-11-25"]).await?, + ["rejected", "sent"] + ); + assert_eq!( + raw_inline_sample_calls(&["2025-11-25", "2026-07-28"]).await?, + ["sent", "rejected"] + ); + Ok(()) +} + +/// A discover session can still pick an older version, where these requests +/// are allowed. +#[tokio::test] +async fn server_requests_allowed_on_older_discover_session() -> anyhow::Result<()> { + for tool in ["sample", "ping"] { + assert_eq!( + call_tool_on_discover_pair(tool, ProtocolVersion::V_2025_11_25).await?, + "sent" + ); + } + Ok(()) +} + +/// The client declares no elicitation modes, so this also pins that the +/// protocol check runs before the capability check. +#[cfg(feature = "elicitation")] +#[tokio::test] +async fn elicit_url_rejected_on_modern_protocol() -> anyhow::Result<()> { + assert_eq!(call_tool_on_modern_pair("elicit_url").await?, "rejected"); + assert_eq!( + call_tool_on_discover_pair("elicit_url", ProtocolVersion::V_2026_07_28).await?, + "rejected" + ); + Ok(()) +} + // A compliant rmcp server cannot produce a server-to-client request at // >= 2026-07-28 (send-side enforcement blocks it), so the client's // receive-side enforcement is exercised with a raw JSON-RPC server.