feat: implement SEP-2260 require server requests to associate with client requests

Author: alexhancock

crates/rmcp/src/service.rs

@@ -149,6 +149,15 @@ pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone {
     ) -> impl Future<Output = ()> + 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<R> =
     JsonRpcMessage<<R as ServiceRole>::Req, <R as ServiceRole>::Resp, <R as ServiceRole>::Not>;
 pub type RxJsonRpcMessage<R> = JsonRpcMessage<
@@ -725,6 +742,11 @@ impl<R: ServiceRole> Peer<R> {
         options: PeerRequestOptions,
         subscription_sender: Option<SubscriptionChannel<R::PeerNot>>,
     ) -> Result<RequestHandle<R>, 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) => {

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.

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<crate::model::RequestId>,
+    future: TaskFuture,
+) -> Result<CallToolResult, TaskExit> {
+    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::<bool>));
+        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::<bool>));
+        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");
+    }
 }

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<Mutex<Option<oneshot::Sender<Result<(), ServiceError>>>>>,
+}
+
+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<RoleServer>,
+    ) -> Result<CallToolResponse, rmcp::ErrorData> {
+        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<RoleClient>,
+    ) -> Result<CreateMessageResult, rmcp::ErrorData> {
+        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(())
+}