feat: SEP-2260 stream-based enforcement of client receive-side request association

crates/rmcp/src/service.rs

@@ -168,12 +168,31 @@ pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone {
     fn enforce_peer_request_association(
         _peer_request: &Self::PeerReq,
         _peer_info: Option<&Self::PeerInfo>,
-        _has_pending_outbound_request: bool,
+        _association: PeerRequestAssociation,
     ) -> Result<(), McpError> {
         Ok(())
     }
 }
 
+/// How an inbound peer request relates to this side's in-flight outbound
+/// requests (SEP-2260).
+///
+/// SEP-2260 defines no wire field for association, so only stream-separating
+/// transports (streamable HTTP) can observe it; other transports yield
+/// [`Self::Unknown`].
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[expect(clippy::exhaustive_enums, reason = "intentionally exhaustive")]
+pub enum PeerRequestAssociation {
+    /// Arrived on the response stream of an in-flight outbound request.
+    Associated,
+    /// Arrived on a stream tied to no in-flight outbound request (e.g. the
+    /// streamable HTTP standalone GET stream).
+    Unassociated,
+    /// The transport cannot distinguish streams; only the coarse in-flight
+    /// signal is available.
+    Unknown { has_pending_outbound_request: bool },
+}
+
 pub(crate) fn uses_legacy_lifecycle(
     protocol_version: Option<&ProtocolVersion>,
     uses_discover_lifecycle: bool,
@@ -182,6 +201,25 @@ pub(crate) fn uses_legacy_lifecycle(
         && protocol_version.is_none_or(|version| version < &ProtocolVersion::V_2026_07_28)
 }
 
+pub(crate) fn peer_request_association<Req: crate::model::GetExtensions, V>(
+    request: &Req,
+    local_responder_pool: &std::collections::HashMap<RequestId, V>,
+) -> PeerRequestAssociation {
+    match request.extensions().get::<InboundStreamOrigin>() {
+        None => PeerRequestAssociation::Unknown {
+            has_pending_outbound_request: !local_responder_pool.is_empty(),
+        },
+        Some(InboundStreamOrigin::Unassociated) => PeerRequestAssociation::Unassociated,
+        Some(InboundStreamOrigin::OutboundRequest(id)) => {
+            if local_responder_pool.contains_key(id) {
+                PeerRequestAssociation::Associated
+            } else {
+                PeerRequestAssociation::Unassociated
+            }
+        }
+    }
+}
+
 tokio::task_local! {
     pub(crate) static ORIGINATING_REQUEST: RequestId;
 }
@@ -205,10 +243,26 @@ pub(crate) fn in_request_handler_scope() -> bool {
 /// outside a handler they return an `invalid_request` error. The association
 /// is task-local and does not cross `tokio::spawn`, so use the task manager
 /// for long-running work.
+///
+/// The client receive-side mirror is [`InboundStreamOrigin`].
 #[derive(Debug, Clone, PartialEq, Eq)]
 #[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")]
 pub struct OriginatingRequestId(pub RequestId);
 
+/// Marker in an inbound request's non-serialized [`Extensions`] recording
+/// which HTTP response stream it arrived on: the receive-side mirror of
+/// [`OriginatingRequestId`]. Never on the wire (SEP-2260 defines no wire
+/// field); when absent, the coarse in-flight check applies.
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[expect(clippy::exhaustive_enums, reason = "intentionally exhaustive")]
+pub enum InboundStreamOrigin {
+    /// The standalone GET stream, or a POST response stream not tied to an
+    /// outbound request.
+    Unassociated,
+    /// The SSE response stream of the POST that carried this outbound request.
+    OutboundRequest(RequestId),
+}
+
 pub type TxJsonRpcMessage<R> =
     JsonRpcMessage<<R as ServiceRole>::Req, <R as ServiceRole>::Resp, <R as ServiceRole>::Not>;
 pub type RxJsonRpcMessage<R> = JsonRpcMessage<
@@ -1448,7 +1502,7 @@ where
                     if let Err(error) = R::enforce_peer_request_association(
                         &request,
                         peer.peer_info().as_deref(),
-                        !local_responder_pool.is_empty(),
+                        peer_request_association(&request, &local_responder_pool),
                     ) {
                         tracing::warn!(%id, message = %error.message, "rejected peer request");
                         // send directly: the sink proxy path would drop the
@@ -1729,4 +1783,73 @@ mod sep2260_marker_tests {
         let request = send_and_capture(None).await;
         assert!(request.extensions().get::<OriginatingRequestId>().is_none());
     }
+
+    #[test]
+    #[expect(
+        deprecated,
+        reason = "Sampling is deprecated by SEP-2577 but remains the canonical restricted request"
+    )]
+    fn peer_request_association_maps_stream_origin() {
+        use std::collections::HashMap;
+
+        use crate::model::{
+            CreateMessageRequest, CreateMessageRequestParams, SamplingMessage, ServerRequest,
+        };
+
+        fn sampling(origin: Option<InboundStreamOrigin>) -> ServerRequest {
+            let mut request = CreateMessageRequest::new(CreateMessageRequestParams::new(
+                vec![SamplingMessage::user_text("hi")],
+                16,
+            ));
+            if let Some(origin) = origin {
+                request.extensions.insert(origin);
+            }
+            ServerRequest::CreateMessageRequest(request)
+        }
+
+        let empty: HashMap<RequestId, ()> = HashMap::new();
+        let in_flight: HashMap<RequestId, ()> = HashMap::from([(RequestId::Number(7), ())]);
+
+        // No marker (stdio): coarse signal.
+        assert_eq!(
+            peer_request_association(&sampling(None), &in_flight),
+            PeerRequestAssociation::Unknown {
+                has_pending_outbound_request: true
+            }
+        );
+        assert_eq!(
+            peer_request_association(&sampling(None), &empty),
+            PeerRequestAssociation::Unknown {
+                has_pending_outbound_request: false
+            }
+        );
+        // Standalone GET stream: unassociated even with requests in flight.
+        assert_eq!(
+            peer_request_association(
+                &sampling(Some(InboundStreamOrigin::Unassociated)),
+                &in_flight
+            ),
+            PeerRequestAssociation::Unassociated
+        );
+        // Originating POST stream of an in-flight request: associated.
+        assert_eq!(
+            peer_request_association(
+                &sampling(Some(InboundStreamOrigin::OutboundRequest(
+                    RequestId::Number(7)
+                ))),
+                &in_flight
+            ),
+            PeerRequestAssociation::Associated
+        );
+        // Stream of a request that is no longer in flight: unassociated.
+        assert_eq!(
+            peer_request_association(
+                &sampling(Some(InboundStreamOrigin::OutboundRequest(
+                    RequestId::Number(8)
+                ))),
+                &in_flight
+            ),
+            PeerRequestAssociation::Unassociated
+        );
+    }
 }

crates/rmcp/src/service/client.rs

@@ -238,14 +238,13 @@ impl ServiceRole for RoleClient {
         }
     }
 
-    // SEP-2260: with no outbound request in flight there is nothing the
-    // server request could be associated with, so reject it. With one in
-    // flight we cannot tell which request it belongs to (no wire field), so
-    // we accept — an under-approximation of the spec's SHOULD.
+    // SEP-2260: reject restricted server requests that arrived unassociated
+    // with any in-flight outbound request. Without stream separation
+    // (`Unknown`) the coarse in-flight check under-approximates the SHOULD.
     fn enforce_peer_request_association(
         peer_request: &Self::PeerReq,
         peer_info: Option<&Self::PeerInfo>,
-        has_pending_outbound_request: bool,
+        association: PeerRequestAssociation,
     ) -> Result<(), ErrorData> {
         let restricted = matches!(
             peer_request,
@@ -258,13 +257,24 @@ impl ServiceRole for RoleClient {
         }
         let strict =
             peer_info.is_some_and(|info| info.protocol_version >= ProtocolVersion::V_2026_07_28);
-        if strict && !has_pending_outbound_request {
-            return Err(ErrorData::invalid_params(
+        if !strict {
+            return Ok(());
+        }
+        let associated = match association {
+            PeerRequestAssociation::Associated => true,
+            PeerRequestAssociation::Unassociated => false,
+            PeerRequestAssociation::Unknown {
+                has_pending_outbound_request,
+            } => has_pending_outbound_request,
+        };
+        if associated {
+            Ok(())
+        } else {
+            Err(ErrorData::invalid_params(
                 "SEP-2260: server-to-client requests must be associated with an in-flight client request",
                 None,
-            ));
+            ))
         }
-        Ok(())
     }
 
     async fn invalidate_response_cache(peer: &Peer<Self>, notification: &Self::PeerNot) {
@@ -2215,3 +2225,72 @@ mod tests {
         assert_eq!(peer.discover(meta).await.unwrap(), expected);
     }
 }
+
+#[cfg(test)]
+mod sep2260_association_tests {
+    use super::*;
+    use crate::{
+        model::{
+            CreateMessageRequest, CreateMessageRequestParams, SamplingMessage, ServerCapabilities,
+        },
+        service::PeerRequestAssociation,
+    };
+
+    fn sampling_request() -> ServerRequest {
+        ServerRequest::CreateMessageRequest(CreateMessageRequest::new(
+            CreateMessageRequestParams::new(vec![SamplingMessage::user_text("hi")], 16),
+        ))
+    }
+
+    fn server_info(version: ProtocolVersion) -> ServerInfo {
+        let mut info = ServerInfo::new(ServerCapabilities::default());
+        info.protocol_version = version;
+        info
+    }
+
+    fn enforce(info: &ServerInfo, association: PeerRequestAssociation) -> Result<(), ErrorData> {
+        RoleClient::enforce_peer_request_association(&sampling_request(), Some(info), association)
+    }
+
+    #[test]
+    fn strict_rejects_unassociated() {
+        let info = server_info(ProtocolVersion::V_2026_07_28);
+        let err = enforce(&info, PeerRequestAssociation::Unassociated).unwrap_err();
+        assert_eq!(err.code, crate::model::ErrorCode::INVALID_PARAMS);
+    }
+
+    #[test]
+    fn strict_accepts_associated() {
+        let info = server_info(ProtocolVersion::V_2026_07_28);
+        assert!(enforce(&info, PeerRequestAssociation::Associated).is_ok());
+    }
+
+    #[test]
+    fn strict_unknown_falls_back_to_coarse_check() {
+        let info = server_info(ProtocolVersion::V_2026_07_28);
+        assert!(
+            enforce(
+                &info,
+                PeerRequestAssociation::Unknown {
+                    has_pending_outbound_request: true
+                }
+            )
+            .is_ok()
+        );
+        assert!(
+            enforce(
+                &info,
+                PeerRequestAssociation::Unknown {
+                    has_pending_outbound_request: false
+                }
+            )
+            .is_err()
+        );
+    }
+
+    #[test]
+    fn legacy_protocol_accepts_even_unassociated() {
+        let info = server_info(ProtocolVersion::V_2025_11_25);
+        assert!(enforce(&info, PeerRequestAssociation::Unassociated).is_ok());
+    }
+}