From d772e4cb0386c4f4820d48c5d74946f2abc26f69 Mon Sep 17 00:00:00 2001 From: Alex Hancock Date: Wed, 22 Jul 2026 14:55:31 -0400 Subject: [PATCH] feat: implement SEP-2260 require server requests to associate with client requests --- crates/rmcp/src/service.rs | 27 ++- crates/rmcp/src/service/server.rs | 25 +++ crates/rmcp/src/task_manager.rs | 80 ++++++++- .../test_sep_2260_request_association.rs | 159 ++++++++++++++++++ 4 files changed, 288 insertions(+), 3 deletions(-) create mode 100644 crates/rmcp/tests/test_sep_2260_request_association.rs diff --git a/crates/rmcp/src/service.rs b/crates/rmcp/src/service.rs index 75f34d30..af5d412f 100644 --- a/crates/rmcp/src/service.rs +++ b/crates/rmcp/src/service.rs @@ -149,6 +149,15 @@ pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone { ) -> impl Future + MaybeSendFuture { async {} } + + #[doc(hidden)] + fn enforce_request_association( + _request: &Self::Req, + _peer_info: Option<&Self::PeerInfo>, + _in_request_handler_scope: bool, + ) -> Result<(), ServiceError> { + Ok(()) + } } pub(crate) fn uses_legacy_lifecycle( @@ -159,6 +168,14 @@ pub(crate) fn uses_legacy_lifecycle( && protocol_version.is_none_or(|version| version < &ProtocolVersion::V_2026_07_28) } +tokio::task_local! { + pub(crate) static ORIGINATING_REQUEST: RequestId; +} + +pub(crate) fn in_request_handler_scope() -> bool { + ORIGINATING_REQUEST.try_with(|_| ()).is_ok() +} + pub type TxJsonRpcMessage = JsonRpcMessage<::Req, ::Resp, ::Not>; pub type RxJsonRpcMessage = JsonRpcMessage< @@ -725,6 +742,11 @@ impl Peer { options: PeerRequestOptions, subscription_sender: Option>, ) -> Result, ServiceError> { + R::enforce_request_association( + &request, + self.peer_info().as_deref(), + in_request_handler_scope(), + )?; let id = self.request_id_provider.next_request_id(); let progress_token = self.progress_token_provider.next_progress_token(); if let Some(metadata) = self.client_request_metadata.get() { @@ -1398,9 +1420,10 @@ where extensions, }; let current_span = tracing::Span::current(); + let handler_id = id.clone(); spawn_service_task(async move { - let result = service - .handle_request(request, context) + let result = ORIGINATING_REQUEST + .scope(handler_id, 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 f4529908..6d2a0d6b 100644 --- a/crates/rmcp/src/service/server.rs +++ b/crates/rmcp/src/service/server.rs @@ -49,6 +49,31 @@ impl ServiceRole for RoleServer { _ => None, } } + + fn enforce_request_association( + request: &Self::Req, + peer_info: Option<&Self::PeerInfo>, + in_request_handler_scope: bool, + ) -> Result<(), ServiceError> { + let restricted = matches!( + request, + ServerRequest::CreateMessageRequest(_) + | ServerRequest::ListRootsRequest(_) + | ServerRequest::ElicitRequest(_) + ); + if !restricted { + return Ok(()); + } + let strict = + peer_info.is_some_and(|info| info.protocol_version >= ProtocolVersion::V_2026_07_28); + if strict && !in_request_handler_scope { + return Err(ServiceError::McpError(ErrorData::invalid_request( + "SEP-2260: server-to-client requests must be associated with an originating client request", + None, + ))); + } + Ok(()) + } } /// It represents the error that may occur when serving the server. diff --git a/crates/rmcp/src/task_manager.rs b/crates/rmcp/src/task_manager.rs index df1c389a..0e8339c0 100644 --- a/crates/rmcp/src/task_manager.rs +++ b/crates/rmcp/src/task_manager.rs @@ -349,8 +349,11 @@ impl TaskManager { let future = make_future(context); let inner = self.inner.clone(); let id_for_task = task_id.clone(); + let originating_request = crate::service::ORIGINATING_REQUEST + .try_with(|id| id.clone()) + .ok(); let handle = tokio::spawn(async move { - let result = future.await; + let result = run_task_operation(originating_request, future).await; let mut inner = inner.lock().expect("task manager lock poisoned"); if let Some(entry) = inner.tasks.get_mut(&id_for_task) { if entry.terminal.is_none() { @@ -538,6 +541,16 @@ fn unknown_task(task_id: &str) -> McpError { McpError::invalid_params(format!("unknown task: {task_id}"), None) } +async fn run_task_operation( + originating_request: Option, + future: TaskFuture, +) -> Result { + match originating_request { + Some(id) => crate::service::ORIGINATING_REQUEST.scope(id, future).await, + None => future.await, + } +} + fn result_to_object(result: &CallToolResult) -> JsonObject { match serde_json::to_value(result) { Ok(serde_json::Value::Object(map)) => map, @@ -917,4 +930,69 @@ mod tests { } panic!("task did not complete after input response"); } + + #[tokio::test] + async fn task_operation_reestablishes_request_association_scope() { + use crate::{ + model::RequestId, + service::{ORIGINATING_REQUEST, in_request_handler_scope}, + }; + + let manager = TaskManager::new(); + let observed = Arc::new(Mutex::new(None::)); + let observed_in_task = observed.clone(); + + ORIGINATING_REQUEST + .scope(RequestId::Number(7), async { + manager.spawn(TaskOptions::default(), move |_ctx| { + let observed_in_task = observed_in_task.clone(); + Box::pin(async move { + *observed_in_task.lock().unwrap() = Some(in_request_handler_scope()); + Ok(ok_result("done")) + }) + }) + }) + .await; + + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + if let Some(scoped) = *observed.lock().unwrap() { + assert!( + scoped, + "task operation must run inside the originating request's association scope" + ); + return; + } + } + panic!("task operation did not run"); + } + + #[tokio::test] + async fn task_operation_without_originating_request_is_unscoped() { + use crate::service::in_request_handler_scope; + + let manager = TaskManager::new(); + let observed = Arc::new(Mutex::new(None::)); + let observed_in_task = observed.clone(); + + manager.spawn(TaskOptions::default(), move |_ctx| { + let observed_in_task = observed_in_task.clone(); + Box::pin(async move { + *observed_in_task.lock().unwrap() = Some(in_request_handler_scope()); + Ok(ok_result("done")) + }) + }); + + for _ in 0..100 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + if let Some(scoped) = *observed.lock().unwrap() { + assert!( + !scoped, + "task operation started without an originating request must remain unscoped" + ); + return; + } + } + panic!("task operation did not run"); + } } diff --git a/crates/rmcp/tests/test_sep_2260_request_association.rs b/crates/rmcp/tests/test_sep_2260_request_association.rs new file mode 100644 index 00000000..de25b8a4 --- /dev/null +++ b/crates/rmcp/tests/test_sep_2260_request_association.rs @@ -0,0 +1,159 @@ +#![cfg(all(feature = "server", feature = "client", not(feature = "local")))] +#![allow(deprecated)] + +use std::sync::{Arc, Mutex}; + +use rmcp::{ + ClientHandler, RoleClient, RoleServer, ServerHandler, ServiceError, ServiceExt, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ClientInfo, ContentBlock, + CreateMessageRequest, CreateMessageRequestParams, CreateMessageResult, ProtocolVersion, + SamplingMessage, ServerCapabilities, ServerInfo, ServerRequest, + }, + service::RequestContext, +}; +use tokio::sync::oneshot; + +#[derive(Clone)] +struct SamplingServer { + outside: Arc>>>>, +} + +impl ServerHandler for SamplingServer { + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + let peer = context.peer.clone(); + let slot = self.outside.clone(); + + let use_generic = request.name == "sample_generic"; + tokio::spawn(async move { + let outside = if use_generic { + peer.send_request(ServerRequest::CreateMessageRequest( + CreateMessageRequest::new(CreateMessageRequestParams::new( + vec![SamplingMessage::user_text("standalone-generic")], + 16, + )), + )) + .await + .map(|_| ()) + } else { + peer.create_message(CreateMessageRequestParams::new( + vec![SamplingMessage::user_text("standalone")], + 16, + )) + .await + .map(|_| ()) + }; + if let Some(tx) = slot.lock().unwrap().take() { + let _ = tx.send(outside); + } + }); + + let nested = context + .peer + .create_message(CreateMessageRequestParams::new( + vec![SamplingMessage::user_text("nested")], + 16, + )) + .await; + nested.map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?; + Ok(CallToolResult::success(vec![ContentBlock::text("ok")]).into()) + } +} + +#[derive(Clone)] +struct SamplingClient; + +impl ClientHandler for SamplingClient { + async fn create_message( + &self, + _params: CreateMessageRequestParams, + _context: RequestContext, + ) -> Result { + Ok(CreateMessageResult::new( + SamplingMessage::assistant_text("pong"), + "test-model".to_string(), + ) + .with_stop_reason(CreateMessageResult::STOP_REASON_END_TURN)) + } + + fn get_info(&self) -> ClientInfo { + let mut info = ClientInfo::default(); + info.protocol_version = ProtocolVersion::V_2026_07_28; + info + } +} + +#[tokio::test] +async fn nested_sampling_allowed_standalone_rejected() -> anyhow::Result<()> { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let (tx, rx) = oneshot::channel(); + let server = SamplingServer { + outside: Arc::new(Mutex::new(Some(tx))), + }; + let server_handle = tokio::spawn(async move { + let running = server.serve(server_transport).await?; + running.waiting().await?; + anyhow::Ok(()) + }); + + let client = SamplingClient.serve(client_transport).await?; + + let result = client + .peer() + .call_tool(CallToolRequestParams::new("sample")) + .await?; + assert_eq!( + result.content.first().unwrap().as_text().unwrap().text, + "ok" + ); + + let outside = rx.await?; + assert!(matches!(outside, Err(ServiceError::McpError(_)))); + + client.cancel().await?; + let _ = server_handle.await?; + Ok(()) +} + +#[tokio::test] +async fn generic_send_request_bypass_rejected() -> anyhow::Result<()> { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let (tx, rx) = oneshot::channel(); + let server = SamplingServer { + outside: Arc::new(Mutex::new(Some(tx))), + }; + let server_handle = tokio::spawn(async move { + let running = server.serve(server_transport).await?; + running.waiting().await?; + anyhow::Ok(()) + }); + + let client = SamplingClient.serve(client_transport).await?; + + let result = client + .peer() + .call_tool(CallToolRequestParams::new("sample_generic")) + .await?; + assert_eq!( + result.content.first().unwrap().as_text().unwrap().text, + "ok" + ); + + let outside = rx.await?; + assert!( + matches!(outside, Err(ServiceError::McpError(_))), + "generic send_request must not bypass SEP-2260 enforcement" + ); + + client.cancel().await?; + let _ = server_handle.await?; + Ok(()) +}