feat: SEP-2260 stream-based enforcement of client receive-side request association (#1055)
* refactor(service): express SEP-2260 receive-side association as an enum Replaces the has_pending_outbound_request bool on enforce_peer_request_association with PeerRequestAssociation, so a stream-separating transport can report per-request association (#1033). Behavior-preserving: the event loop still passes the coarse signal as Unknown. * feat(transport): record inbound stream origin on streamable HTTP client (#1033) The worker attaches an InboundStreamOrigin extension to each inbound server request: Unassociated for the standalone GET stream, OutboundRequest(id) for a POST's SSE stream. Mirror of the OriginatingRequestId marker used by the server side. * feat(service): enforce SEP-2260 client receive-side check per stream origin (#1033) The event loop maps InboundStreamOrigin plus the in-flight responder pool to PeerRequestAssociation: restricted requests arriving on the standalone GET stream are now rejected with -32602 even while unrelated outbound requests are in flight. * test: end-to-end SEP-2260 stream-based enforcement over streamable HTTP (#1033) * docs: cross-link SEP-2260 stream markers * test: origin marker survives SSE resumption per SEP-2260 (#1033) A POST SSE stream resumed via GET + Last-Event-ID (SEP-1699) reconnects beneath execute_sse_stream, so requests replayed after a resume keep their OutboundRequest origin. Pins the layering invariant: hoisting reconnection above the marker attach point would wrongly reject associated requests with -32602. * docs: trim SEP-2260 comments to spec rationale * docs: align SEP-2577 expect reason with sibling suppressions * test: harden SEP-2260 stream-enforcement e2e (#1033) Detect handler invocation via a channel asserted empty instead of a panic in a spawned task (swallowed, cannot fail the test); surface scripted-server misuse as transport errors rather than panics in the transport task; bound the tail awaits with 5s timeouts. Correct the header comment: a 2026-07-28 server minting a session id is not spec-legal (SEP-2567 removes sessions and the GET endpoint) — the scripted server is deliberately non-conforming, which is the point of receive-side enforcement. * test: adapt SEP-2260 association tests to ServerPeerInfo Rebase onto 3.0.0: RoleClient::PeerInfo is now ServerPeerInfo (#1065) and its constructor takes the protocol version directly. * chore: add 'Copy' macro to new PeerRequestAssociation Co-authored-by: Dale Seo <5466341+DaleSeo@users.noreply.github.com> --------- Co-authored-by: Dale Seo <5466341+DaleSeo@users.noreply.github.com>
This commit is contained in:
camillelawrence
GitHub
parent
3bb3c8dbce
commit
1d0473ba75
4 changed files with 734 additions and 16 deletions
|
|
@ -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, Copy, 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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
@ -2209,3 +2219,73 @@ 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) -> ServerPeerInfo {
|
||||
ServerPeerInfo::new(version, ServerCapabilities::default())
|
||||
}
|
||||
|
||||
fn enforce(
|
||||
info: &ServerPeerInfo,
|
||||
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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -19,10 +19,11 @@ use super::common::client_side_sse::{
|
|||
use crate::{
|
||||
RoleClient,
|
||||
model::{
|
||||
ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData, GetMeta,
|
||||
ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData, GetExtensions, GetMeta,
|
||||
InitializedNotification, JsonObject, ProtocolVersion, RequestId, ServerJsonRpcMessage,
|
||||
ServerResult,
|
||||
},
|
||||
service::InboundStreamOrigin,
|
||||
transport::{
|
||||
common::{client_side_sse::SseAutoReconnectStream, mcp_headers},
|
||||
worker::{Worker, WorkerQuitReason, WorkerSendRequest, WorkerTransport},
|
||||
|
|
@ -635,6 +636,7 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
|||
+ Send
|
||||
+ 'static,
|
||||
sse_worker_tx: tokio::sync::mpsc::Sender<ServerJsonRpcMessage>,
|
||||
origin: InboundStreamOrigin,
|
||||
close_on_response: bool,
|
||||
ct: CancellationToken,
|
||||
) -> Result<(), StreamableHttpError<C::Error>> {
|
||||
|
|
@ -649,9 +651,14 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
|||
break;
|
||||
}
|
||||
};
|
||||
let Some(message) = message.transpose()? else {
|
||||
let Some(mut message) = message.transpose()? else {
|
||||
break;
|
||||
};
|
||||
// SEP-2260: mark inbound requests with the stream they arrived on
|
||||
// for the client receive-side association check.
|
||||
if let ServerJsonRpcMessage::Request(request) = &mut message {
|
||||
request.request.extensions_mut().insert(origin.clone());
|
||||
}
|
||||
let is_response = matches!(
|
||||
message,
|
||||
ServerJsonRpcMessage::Response(_) | ServerJsonRpcMessage::Error(_)
|
||||
|
|
@ -719,6 +726,7 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
|||
Self::execute_sse_stream(
|
||||
sse_stream,
|
||||
sse_worker_tx,
|
||||
InboundStreamOrigin::Unassociated,
|
||||
false,
|
||||
transport_task_ct.child_token(),
|
||||
)
|
||||
|
|
@ -1283,9 +1291,18 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
);
|
||||
}
|
||||
let stream_tx = sse_worker_tx.clone();
|
||||
let origin = match &stream_request_id {
|
||||
Some(id) => {
|
||||
InboundStreamOrigin::OutboundRequest(
|
||||
id.clone(),
|
||||
)
|
||||
}
|
||||
None => InboundStreamOrigin::Unassociated,
|
||||
};
|
||||
streams.spawn(async move {
|
||||
let result = Self::execute_sse_stream(
|
||||
sse_stream, stream_tx, true, stream_ct,
|
||||
sse_stream, stream_tx, origin, true,
|
||||
stream_ct,
|
||||
)
|
||||
.await;
|
||||
(stream_request_id, result)
|
||||
|
|
@ -1342,9 +1359,13 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
.insert(request_id.clone(), stream_ct.clone());
|
||||
}
|
||||
let stream_tx = sse_worker_tx.clone();
|
||||
let origin = match &stream_request_id {
|
||||
Some(id) => InboundStreamOrigin::OutboundRequest(id.clone()),
|
||||
None => InboundStreamOrigin::Unassociated,
|
||||
};
|
||||
streams.spawn(async move {
|
||||
let result = Self::execute_sse_stream(
|
||||
sse_stream, stream_tx, true, stream_ct,
|
||||
sse_stream, stream_tx, origin, true, stream_ct,
|
||||
)
|
||||
.await;
|
||||
(stream_request_id, result)
|
||||
|
|
@ -1769,7 +1790,66 @@ mod tests {
|
|||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::model::{ListToolsResult, NumberOrString, ServerResult, Tool};
|
||||
use crate::{
|
||||
model::{
|
||||
GetExtensions, ListToolsResult, NumberOrString, ServerRequest, ServerResult, Tool,
|
||||
},
|
||||
service::InboundStreamOrigin,
|
||||
};
|
||||
|
||||
#[expect(
|
||||
deprecated,
|
||||
reason = "Sampling is deprecated by SEP-2577 but remains the canonical restricted request"
|
||||
)]
|
||||
fn sampling_request_message(id: i64) -> ServerJsonRpcMessage {
|
||||
use crate::model::{CreateMessageRequest, CreateMessageRequestParams, SamplingMessage};
|
||||
ServerJsonRpcMessage::request(
|
||||
ServerRequest::CreateMessageRequest(CreateMessageRequest::new(
|
||||
CreateMessageRequestParams::new(vec![SamplingMessage::user_text("hi")], 16),
|
||||
)),
|
||||
NumberOrString::Number(id),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_sse_stream_marks_inbound_requests_with_origin() {
|
||||
for origin in [
|
||||
InboundStreamOrigin::Unassociated,
|
||||
InboundStreamOrigin::OutboundRequest(RequestId::Number(3)),
|
||||
] {
|
||||
let response = ServerJsonRpcMessage::response(
|
||||
ServerResult::ListToolsResult(ListToolsResult::default()),
|
||||
NumberOrString::Number(1),
|
||||
);
|
||||
let stream = futures::stream::iter([Ok(sampling_request_message(9)), Ok(response)]);
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel(4);
|
||||
StreamableHttpClientWorker::<StatelessReconnectClient>::execute_sse_stream(
|
||||
stream,
|
||||
tx,
|
||||
origin.clone(),
|
||||
false,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
.await
|
||||
.expect("stream completes");
|
||||
|
||||
let ServerJsonRpcMessage::Request(request) =
|
||||
rx.recv().await.expect("request forwarded")
|
||||
else {
|
||||
panic!("expected request first");
|
||||
};
|
||||
assert_eq!(
|
||||
request.request.extensions().get::<InboundStreamOrigin>(),
|
||||
Some(&origin),
|
||||
"inbound requests must carry their stream origin"
|
||||
);
|
||||
// Responses are correlated by JSON-RPC id; no marker needed or added.
|
||||
assert!(matches!(
|
||||
rx.recv().await.expect("response forwarded"),
|
||||
ServerJsonRpcMessage::Response(_)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
type ReconnectAttempt = (Option<String>, Option<String>);
|
||||
|
||||
|
|
@ -1868,6 +1948,132 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct ResumedRequestClient {
|
||||
reconnects: Arc<Mutex<Vec<ReconnectAttempt>>>,
|
||||
}
|
||||
|
||||
impl StreamableHttpClient for ResumedRequestClient {
|
||||
type Error = std::io::Error;
|
||||
|
||||
async fn post_message(
|
||||
&self,
|
||||
_uri: Arc<str>,
|
||||
_message: ClientJsonRpcMessage,
|
||||
_session_id: Option<Arc<str>>,
|
||||
_auth_header: Option<String>,
|
||||
_custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||
Err(StreamableHttpError::UnexpectedServerResponse(
|
||||
"unexpected POST".into(),
|
||||
))
|
||||
}
|
||||
|
||||
async fn delete_session(
|
||||
&self,
|
||||
_uri: Arc<str>,
|
||||
_session_id: Arc<str>,
|
||||
_auth_header: Option<String>,
|
||||
_custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<(), StreamableHttpError<Self::Error>> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_stream(
|
||||
&self,
|
||||
_uri: Arc<str>,
|
||||
session_id: Option<Arc<str>>,
|
||||
last_event_id: Option<String>,
|
||||
_auth_header: Option<String>,
|
||||
_custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<BoxedSseStream, StreamableHttpError<Self::Error>> {
|
||||
self.reconnects
|
||||
.lock()
|
||||
.expect("lock reconnects")
|
||||
.push((session_id.map(|id| id.to_string()), last_event_id));
|
||||
let request = sampling_request_message(9);
|
||||
let response = ServerJsonRpcMessage::response(
|
||||
ServerResult::ListToolsResult(ListToolsResult::default()),
|
||||
NumberOrString::Number(1),
|
||||
);
|
||||
// Stay open after the response, like a live connection, so the
|
||||
// post-response drain in `execute_sse_stream` doesn't trigger
|
||||
// further reconnects.
|
||||
Ok(futures::stream::iter([request, response].map(|message| {
|
||||
Ok(Sse {
|
||||
event: None,
|
||||
data: Some(serde_json::to_string(&message).expect("serialize message")),
|
||||
id: None,
|
||||
retry: None,
|
||||
})
|
||||
}))
|
||||
.chain(futures::stream::pending())
|
||||
.boxed())
|
||||
}
|
||||
}
|
||||
|
||||
/// SEP-1699 resumes a broken POST SSE stream via GET + Last-Event-ID
|
||||
/// beneath `execute_sse_stream`, so the SEP-2260 origin marker must span
|
||||
/// resumes; if reconnection were hoisted above the marker attach point,
|
||||
/// replayed associated requests would be wrongly rejected with -32602.
|
||||
#[tokio::test]
|
||||
async fn resumed_post_stream_requests_keep_outbound_origin() {
|
||||
let initial = futures::stream::iter([Ok(Sse {
|
||||
event: None,
|
||||
data: None,
|
||||
id: Some("e1".into()),
|
||||
retry: Some(0),
|
||||
})])
|
||||
.boxed();
|
||||
let client = ResumedRequestClient::default();
|
||||
let reconnects = client.reconnects.clone();
|
||||
let sse_stream =
|
||||
StreamableHttpClientWorker::<ResumedRequestClient>::response_sse_to_jsonrpc(
|
||||
initial,
|
||||
None,
|
||||
client,
|
||||
Arc::from("http://localhost/mcp"),
|
||||
None,
|
||||
HashMap::new(),
|
||||
DEFAULT_MAX_SSE_EVENT_SIZE,
|
||||
Arc::new(ExponentialBackoff {
|
||||
max_times: Some(1),
|
||||
base_duration: Duration::ZERO,
|
||||
}),
|
||||
);
|
||||
|
||||
let origin = InboundStreamOrigin::OutboundRequest(RequestId::Number(3));
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel(4);
|
||||
StreamableHttpClientWorker::<ResumedRequestClient>::execute_sse_stream(
|
||||
sse_stream,
|
||||
tx,
|
||||
origin.clone(),
|
||||
true,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
.await
|
||||
.expect("stream completes");
|
||||
|
||||
assert_eq!(
|
||||
reconnects.lock().expect("lock reconnects").as_slice(),
|
||||
&[(None, Some("e1".into()))],
|
||||
"the request must arrive on the resumed connection"
|
||||
);
|
||||
let ServerJsonRpcMessage::Request(request) = rx.recv().await.expect("request forwarded")
|
||||
else {
|
||||
panic!("expected request first");
|
||||
};
|
||||
assert_eq!(
|
||||
request.request.extensions().get::<InboundStreamOrigin>(),
|
||||
Some(&origin),
|
||||
"origin marker must survive SSE resumption"
|
||||
);
|
||||
assert!(matches!(
|
||||
rx.recv().await.expect("response forwarded"),
|
||||
ServerJsonRpcMessage::Response(_)
|
||||
));
|
||||
}
|
||||
|
||||
fn tool(name: &'static str, annotation: serde_json::Value) -> Tool {
|
||||
let schema = json!({
|
||||
"type": "object",
|
||||
|
|
|
|||
309
crates/rmcp/tests/test_sep_2260_stream_enforcement.rs
Normal file
309
crates/rmcp/tests/test_sep_2260_stream_enforcement.rs
Normal file
|
|
@ -0,0 +1,309 @@
|
|||
//! SEP-2260 follow-up (#1033): stream-based receive-side enforcement.
|
||||
//!
|
||||
//! Scripted streamable HTTP "server": answers a legacy initialize with
|
||||
//! protocol 2026-07-28 AND a session id. That combination is NOT
|
||||
//! spec-compliant: the 2026-07-28 revision removes protocol-level sessions
|
||||
//! and the standalone GET stream (SEP-2567; transports spec: "do not mint
|
||||
//! or echo session IDs"). Receive-side enforcement (#1033) exists precisely
|
||||
//! to protect the client from non-conforming servers, and rmcp's client
|
||||
//! tolerates the session id and opens the standalone GET stream — so this
|
||||
//! is the reachable path where the client has BOTH a GET stream and strict
|
||||
//! SEP-2260 enforcement.
|
||||
#![cfg(all(
|
||||
feature = "client",
|
||||
feature = "transport-streamable-http-client",
|
||||
not(feature = "local")
|
||||
))]
|
||||
#![expect(
|
||||
deprecated,
|
||||
reason = "Sampling is deprecated by SEP-2577 but remains the canonical restricted request"
|
||||
)]
|
||||
|
||||
use std::{collections::HashMap, sync::Arc, time::Duration};
|
||||
|
||||
use futures::{StreamExt, stream::BoxStream};
|
||||
use http::{HeaderName, HeaderValue};
|
||||
use rmcp::{
|
||||
ClientHandler,
|
||||
model::{
|
||||
ClientInfo, ClientJsonRpcMessage, CreateMessageRequestParams, CreateMessageResult,
|
||||
ProtocolVersion, SamplingMessage, ServerCapabilities, ServerInfo, ServerJsonRpcMessage,
|
||||
},
|
||||
service::{ClientLifecycleMode, RequestContext, RoleClient, serve_client_with_lifecycle},
|
||||
transport::streamable_http_client::{
|
||||
StreamableHttpClient, StreamableHttpClientTransport, StreamableHttpClientTransportConfig,
|
||||
StreamableHttpError, StreamableHttpPostResponse,
|
||||
},
|
||||
};
|
||||
use serde_json::{Value, json};
|
||||
use sse_stream::{Error as SseError, Sse};
|
||||
use tokio::sync::{Mutex, mpsc};
|
||||
|
||||
fn to_sse(message: Value) -> Result<Sse, SseError> {
|
||||
Ok(Sse {
|
||||
event: None,
|
||||
data: Some(message.to_string()),
|
||||
id: None,
|
||||
retry: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn message_stream(rx: mpsc::Receiver<Value>) -> BoxStream<'static, Result<Sse, SseError>> {
|
||||
tokio_stream::wrappers::ReceiverStream::new(rx)
|
||||
.map(to_sse)
|
||||
.boxed()
|
||||
}
|
||||
|
||||
/// Scripted server: initialize -> JSON init result (2026-07-28 + session);
|
||||
/// first non-initialize request POST -> SSE stream fed by `post_stream`;
|
||||
/// everything else -> Accepted. Every message the client POSTs is forwarded
|
||||
/// to `posted`.
|
||||
#[derive(Clone)]
|
||||
struct ScriptedServer {
|
||||
get_stream: Arc<Mutex<Option<mpsc::Receiver<Value>>>>,
|
||||
post_stream: Arc<Mutex<Option<mpsc::Receiver<Value>>>>,
|
||||
posted: mpsc::UnboundedSender<Value>,
|
||||
}
|
||||
|
||||
impl StreamableHttpClient for ScriptedServer {
|
||||
type Error = std::io::Error;
|
||||
|
||||
async fn post_message(
|
||||
&self,
|
||||
_uri: Arc<str>,
|
||||
message: ClientJsonRpcMessage,
|
||||
_session_id: Option<Arc<str>>,
|
||||
_auth_header: Option<String>,
|
||||
_custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||
let value = serde_json::to_value(&message).expect("serialize client message");
|
||||
// Receiver drop is normal at test teardown; never panic in the transport task.
|
||||
let _ = self.posted.send(value.clone());
|
||||
if value["method"] == "initialize" {
|
||||
let mut info = ServerInfo::new(ServerCapabilities::default());
|
||||
info.protocol_version = ProtocolVersion::V_2026_07_28;
|
||||
let response = ServerJsonRpcMessage::response(
|
||||
rmcp::model::ServerResult::InitializeResult(info),
|
||||
serde_json::from_value(value["id"].clone()).expect("request id"),
|
||||
);
|
||||
return Ok(StreamableHttpPostResponse::Json(
|
||||
response,
|
||||
Some("scripted-session".into()),
|
||||
));
|
||||
}
|
||||
if matches!(message, ClientJsonRpcMessage::Request(_)) {
|
||||
// Fail as a transport error rather than panicking: this code runs
|
||||
// in the transport task, where a panic is swallowed and shows up
|
||||
// only as an opaque timeout in the test.
|
||||
let rx = self.post_stream.lock().await.take().ok_or_else(|| {
|
||||
StreamableHttpError::Client(std::io::Error::other(
|
||||
"scripted server expects exactly one non-initialize request POST",
|
||||
))
|
||||
})?;
|
||||
return Ok(StreamableHttpPostResponse::Sse(message_stream(rx), None));
|
||||
}
|
||||
Ok(StreamableHttpPostResponse::Accepted)
|
||||
}
|
||||
|
||||
async fn delete_session(
|
||||
&self,
|
||||
_uri: Arc<str>,
|
||||
_session_id: Arc<str>,
|
||||
_auth_header: Option<String>,
|
||||
_custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<(), StreamableHttpError<Self::Error>> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_stream(
|
||||
&self,
|
||||
_uri: Arc<str>,
|
||||
_session_id: Option<Arc<str>>,
|
||||
_last_event_id: Option<String>,
|
||||
_auth_header: Option<String>,
|
||||
_custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<BoxStream<'static, Result<Sse, SseError>>, StreamableHttpError<Self::Error>> {
|
||||
match self.get_stream.lock().await.take() {
|
||||
Some(rx) => Ok(message_stream(rx)),
|
||||
// Reconnect after the scripted stream ends: stay silent.
|
||||
None => Ok(futures::stream::pending().boxed()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct SamplingClient {
|
||||
// Every invocation is recorded here. Tests assert on the receiver (the
|
||||
// negative test asserts it stays empty); a panic in this handler would
|
||||
// run in a spawned task and be silently swallowed.
|
||||
sampled: mpsc::UnboundedSender<CreateMessageRequestParams>,
|
||||
}
|
||||
|
||||
impl ClientHandler for SamplingClient {
|
||||
async fn create_message(
|
||||
&self,
|
||||
params: CreateMessageRequestParams,
|
||||
_context: RequestContext<RoleClient>,
|
||||
) -> Result<CreateMessageResult, rmcp::ErrorData> {
|
||||
let _ = self.sampled.send(params);
|
||||
Ok(CreateMessageResult::new(
|
||||
SamplingMessage::assistant_text("pong"),
|
||||
"test-model".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
fn get_info(&self) -> ClientInfo {
|
||||
ClientInfo::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn sampling_request(id: u32) -> Value {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"method": "sampling/createMessage",
|
||||
"params": {
|
||||
"messages": [{ "role": "user", "content": { "type": "text", "text": "hi" } }],
|
||||
"maxTokens": 16
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn tools_list_response(id: &Value) -> Value {
|
||||
json!({ "jsonrpc": "2.0", "id": id, "result": { "tools": [] } })
|
||||
}
|
||||
|
||||
async fn next_posted(posted: &mut mpsc::UnboundedReceiver<Value>) -> Value {
|
||||
tokio::time::timeout(Duration::from_secs(5), posted.recv())
|
||||
.await
|
||||
.expect("posted message within 5s")
|
||||
.expect("channel open")
|
||||
}
|
||||
|
||||
struct Harness {
|
||||
client: rmcp::service::RunningService<RoleClient, SamplingClient>,
|
||||
/// Every message the client POSTs to the scripted server.
|
||||
posted: mpsc::UnboundedReceiver<Value>,
|
||||
/// Feeds the standalone GET stream.
|
||||
get_tx: mpsc::Sender<Value>,
|
||||
/// Feeds the SSE stream of the in-flight tools/list POST.
|
||||
post_tx: mpsc::Sender<Value>,
|
||||
/// In-flight tools/list call (response withheld until the test releases it).
|
||||
call: tokio::task::JoinHandle<Result<rmcp::model::ListToolsResult, rmcp::ServiceError>>,
|
||||
tools_list_id: Value,
|
||||
/// Sampling params seen by the client handler.
|
||||
sampled: mpsc::UnboundedReceiver<CreateMessageRequestParams>,
|
||||
}
|
||||
|
||||
/// Drive startup + one in-flight tools/list.
|
||||
async fn setup() -> Harness {
|
||||
let (get_tx, get_rx) = mpsc::channel(8);
|
||||
let (post_tx, post_rx) = mpsc::channel(8);
|
||||
let (posted_tx, mut posted_rx) = mpsc::unbounded_channel();
|
||||
let (sampled_tx, sampled_rx) = mpsc::unbounded_channel();
|
||||
let server = ScriptedServer {
|
||||
get_stream: Arc::new(Mutex::new(Some(get_rx))),
|
||||
post_stream: Arc::new(Mutex::new(Some(post_rx))),
|
||||
posted: posted_tx,
|
||||
};
|
||||
let transport = StreamableHttpClientTransport::with_client(
|
||||
server,
|
||||
StreamableHttpClientTransportConfig::with_uri("http://scripted/mcp"),
|
||||
);
|
||||
let client = serve_client_with_lifecycle(
|
||||
SamplingClient {
|
||||
sampled: sampled_tx,
|
||||
},
|
||||
transport,
|
||||
ClientLifecycleMode::Initialize,
|
||||
)
|
||||
.await
|
||||
.expect("initialize against scripted server");
|
||||
|
||||
// initialize + notifications/initialized already posted during startup.
|
||||
assert_eq!(next_posted(&mut posted_rx).await["method"], "initialize");
|
||||
assert_eq!(
|
||||
next_posted(&mut posted_rx).await["method"],
|
||||
"notifications/initialized"
|
||||
);
|
||||
|
||||
// Unrelated outbound request, kept in flight (response withheld).
|
||||
let peer = client.peer().clone();
|
||||
let call = tokio::spawn(async move { peer.list_tools(None).await });
|
||||
let tools_list = next_posted(&mut posted_rx).await;
|
||||
assert_eq!(tools_list["method"], "tools/list");
|
||||
let tools_list_id = tools_list["id"].clone();
|
||||
|
||||
Harness {
|
||||
client,
|
||||
posted: posted_rx,
|
||||
get_tx,
|
||||
post_tx,
|
||||
call,
|
||||
tools_list_id,
|
||||
sampled: sampled_rx,
|
||||
}
|
||||
}
|
||||
|
||||
/// #1033 scenario 1: a restricted request on the standalone GET stream while
|
||||
/// an unrelated outbound request is in flight must be rejected with -32602.
|
||||
/// (The coarse check from #1029 incorrectly accepted this.)
|
||||
#[tokio::test]
|
||||
async fn restricted_request_on_get_stream_rejected_while_unrelated_request_in_flight()
|
||||
-> anyhow::Result<()> {
|
||||
let mut h = setup().await;
|
||||
|
||||
h.get_tx.send(sampling_request(100)).await?;
|
||||
|
||||
let rejection = next_posted(&mut h.posted).await;
|
||||
assert_eq!(
|
||||
rejection["id"], 100,
|
||||
"reply to the sampling request: {rejection}"
|
||||
);
|
||||
assert_eq!(
|
||||
rejection["error"]["code"], -32602,
|
||||
"SEP-2260: GET-stream request must be rejected even with an unrelated \
|
||||
request in flight, got {rejection}"
|
||||
);
|
||||
|
||||
h.post_tx
|
||||
.send(tools_list_response(&h.tools_list_id))
|
||||
.await?;
|
||||
tokio::time::timeout(Duration::from_secs(5), h.call).await???;
|
||||
// Rejected means rejected: the handler must not ALSO have been invoked.
|
||||
assert!(
|
||||
h.sampled.try_recv().is_err(),
|
||||
"handler must not see the rejected sampling request"
|
||||
);
|
||||
tokio::time::timeout(Duration::from_secs(5), h.client.cancel()).await??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Positive twin: the same restricted request arriving on the SSE stream of
|
||||
/// the originating POST is dispatched to the handler and answered.
|
||||
#[tokio::test]
|
||||
async fn restricted_request_on_originating_post_stream_is_dispatched() -> anyhow::Result<()> {
|
||||
let mut h = setup().await;
|
||||
|
||||
h.post_tx.send(sampling_request(200)).await?;
|
||||
|
||||
let response = next_posted(&mut h.posted).await;
|
||||
assert_eq!(
|
||||
response["id"], 200,
|
||||
"reply to the sampling request: {response}"
|
||||
);
|
||||
assert_eq!(
|
||||
response["result"]["model"], "test-model",
|
||||
"request on the originating POST stream must reach the handler, got {response}"
|
||||
);
|
||||
tokio::time::timeout(Duration::from_secs(5), h.sampled.recv())
|
||||
.await?
|
||||
.expect("handler invoked");
|
||||
|
||||
h.post_tx
|
||||
.send(tools_list_response(&h.tools_list_id))
|
||||
.await?;
|
||||
tokio::time::timeout(Duration::from_secs(5), h.call).await???;
|
||||
tokio::time::timeout(Duration::from_secs(5), h.client.cancel()).await??;
|
||||
Ok(())
|
||||
}
|
||||
Loading…
Reference in a new issue