feat: route SEP-2260 associated server requests to the originating SSE stream (#1029)
* feat: implement SEP-2260 require server requests to associate with client requests * feat(service): attach originating request id to in-handler outbound requests (SEP-2260) Reuses the ORIGINATING_REQUEST task-local from #1027; the marker rides the request's non-serialized Extensions so the streamable HTTP server can route associated requests to the originating SSE stream. * feat(transport): route associated server requests to originating SSE stream (SEP-2260) * test: end-to-end SEP-2260 stream routing over streamable HTTP * test: register SEP-2260 routing test with required-features; drop dead list_tools * docs: document SEP-2260 association requirement and stream routing * docs: fix redundant explicit link on OriginatingRequestId * docs: note marker is role-agnostic; add reason to deprecated expect Review follow-ups: the OriginatingRequestId rustdoc now states the marker is attached for both roles with only the server transport reading it, and the test's deprecation suppression carries a reason per house style. * docs: condense inline comments to repo convention Trims multi-line inline comments added in this branch down to the terse one-to-two-line style used elsewhere in the codebase, keeping only the non-obvious rationale (spec references, fallback invariant, version choice). * docs: condense SEP-2260 rustdoc sections Tightens the duplicated per-method association section from eleven lines to five, and trims the marker and SessionManager docs to the same information in fewer words. * docs: single-source SEP-2260 caller requirements on OriginatingRequestId * feat(service): reject unassociated server-to-client requests on the client (SEP-2260) Implements the client receive-side of SEP-2260: restricted server requests (sampling/createMessage, roots/list, elicitation/create) received while the client has no outbound request in flight are answered with -32602 invalid params instead of being dispatched to the handler. Gated on negotiated protocol >= 2026-07-28; ping is exempt. With a request in flight we cannot tell which one the server request belongs to (SEP-2260 defines no wire field), so this is a deliberate under-approximation of the spec's SHOULD; exact stream-based enforcement for streamable HTTP needs receive-side provenance plumbing and is left as a follow-up. --------- Co-authored-by: Alex Hancock <alexhancock@block.xyz>
This commit is contained in:
camillelawrence
GitHub
parent
aad4d4e4c6
commit
397d416fdc
9 changed files with 853 additions and 6 deletions
|
|
@ -445,3 +445,8 @@ required-features = [
|
|||
"reqwest",
|
||||
]
|
||||
path = "tests/test_streamable_http_disconnect_cancel.rs"
|
||||
|
||||
[[test]]
|
||||
name = "test_sep_2260_stream_routing"
|
||||
required-features = ["server", "elicitation", "transport-streamable-http-server", "reqwest"]
|
||||
path = "tests/test_sep_2260_stream_routing.rs"
|
||||
|
|
|
|||
|
|
@ -149,6 +149,29 @@ 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(())
|
||||
}
|
||||
|
||||
/// Receive-side counterpart of [`Self::enforce_request_association`]:
|
||||
/// SEP-2260 says clients receiving a server-to-client request with no
|
||||
/// associated outbound request should reject it with invalid params. An
|
||||
/// error return is sent back to the peer instead of dispatching to the
|
||||
/// handler.
|
||||
#[doc(hidden)]
|
||||
fn enforce_peer_request_association(
|
||||
_peer_request: &Self::PeerReq,
|
||||
_peer_info: Option<&Self::PeerInfo>,
|
||||
_has_pending_outbound_request: bool,
|
||||
) -> Result<(), McpError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn uses_legacy_lifecycle(
|
||||
|
|
@ -159,6 +182,33 @@ 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()
|
||||
}
|
||||
|
||||
/// Marker in an outbound request's non-serialized [`Extensions`] identifying
|
||||
/// the in-flight peer request it was issued from (SEP-2260). Attached for both
|
||||
/// roles whenever a request is sent from within a request handler; the
|
||||
/// streamable HTTP server reads it to deliver server-initiated requests on the
|
||||
/// originating request's SSE stream. Never on the wire (SEP-2260 defines no
|
||||
/// wire field), so session managers that serialize messages between processes
|
||||
/// lose it and such requests fall back to the standalone stream with a warning.
|
||||
///
|
||||
/// # Caller requirements
|
||||
///
|
||||
/// From protocol version `2026-07-28`, server-to-client sampling, roots, and
|
||||
/// elicitation requests must be issued while handling a client request;
|
||||
/// 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.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")]
|
||||
pub struct OriginatingRequestId(pub RequestId);
|
||||
|
||||
pub type TxJsonRpcMessage<R> =
|
||||
JsonRpcMessage<<R as ServiceRole>::Req, <R as ServiceRole>::Resp, <R as ServiceRole>::Not>;
|
||||
pub type RxJsonRpcMessage<R> = JsonRpcMessage<
|
||||
|
|
@ -725,6 +775,16 @@ 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(),
|
||||
)?;
|
||||
if let Ok(originating) = ORIGINATING_REQUEST.try_with(|id| id.clone()) {
|
||||
request
|
||||
.extensions_mut()
|
||||
.insert(OriginatingRequestId(originating));
|
||||
}
|
||||
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() {
|
||||
|
|
@ -1389,6 +1449,24 @@ where
|
|||
..
|
||||
})) => {
|
||||
tracing::debug!(%id, ?request, "received request");
|
||||
if let Err(error) = R::enforce_peer_request_association(
|
||||
&request,
|
||||
peer.peer_info().as_deref(),
|
||||
!local_responder_pool.is_empty(),
|
||||
) {
|
||||
tracing::warn!(%id, message = %error.message, "rejected peer request");
|
||||
// send directly: the sink proxy path would drop the
|
||||
// error since the request was never registered in
|
||||
// local_ct_pool
|
||||
let send = transport.send(JsonRpcMessage::error(error, Some(id)));
|
||||
let current_span = tracing::Span::current();
|
||||
response_send_tasks.spawn(async move {
|
||||
if let Err(error) = send.await {
|
||||
tracing::error!(%error, "fail to send rejection error");
|
||||
}
|
||||
}.instrument(current_span));
|
||||
continue;
|
||||
}
|
||||
{
|
||||
let service = shared_service.clone();
|
||||
let sink = sink_proxy_tx.clone();
|
||||
|
|
@ -1409,9 +1487,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) => {
|
||||
|
|
@ -1608,3 +1687,50 @@ where
|
|||
dg: ct.drop_guard(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(test, feature = "server"))]
|
||||
mod sep2260_marker_tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::*;
|
||||
use crate::model::{PingRequest, RequestId, ServerRequest};
|
||||
|
||||
fn ping() -> ServerRequest {
|
||||
ServerRequest::PingRequest(PingRequest {
|
||||
method: Default::default(),
|
||||
extensions: Default::default(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn send_and_capture(scope: Option<RequestId>) -> <RoleServer as ServiceRole>::Req {
|
||||
// peer_info None keeps enforcement non-strict; only the sink message matters.
|
||||
let (peer, mut rx) =
|
||||
Peer::<RoleServer>::new(Arc::new(AtomicU32RequestIdProvider::default()), None);
|
||||
let send = peer.send_request_with_option(ping(), PeerRequestOptions::no_options());
|
||||
let _handle = match scope {
|
||||
Some(id) => ORIGINATING_REQUEST.scope(id, send).await.unwrap(),
|
||||
None => send.await.unwrap(),
|
||||
};
|
||||
let PeerSinkMessage::Request { request, .. } = rx.recv().await.expect("sink message")
|
||||
else {
|
||||
panic!("expected a request sink message");
|
||||
};
|
||||
request
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn outbound_request_carries_originating_id_when_in_scope() {
|
||||
let request = send_and_capture(Some(RequestId::Number(7))).await;
|
||||
let marker = request
|
||||
.extensions()
|
||||
.get::<OriginatingRequestId>()
|
||||
.expect("marker attached");
|
||||
assert_eq!(marker.0, RequestId::Number(7));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn outbound_request_has_no_marker_outside_scope() {
|
||||
let request = send_and_capture(None).await;
|
||||
assert!(request.extensions().get::<OriginatingRequestId>().is_none());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -213,6 +213,35 @@ 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.
|
||||
fn enforce_peer_request_association(
|
||||
peer_request: &Self::PeerReq,
|
||||
peer_info: Option<&Self::PeerInfo>,
|
||||
has_pending_outbound_request: bool,
|
||||
) -> Result<(), ErrorData> {
|
||||
let restricted = matches!(
|
||||
peer_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 && !has_pending_outbound_request {
|
||||
return 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) {
|
||||
match notification {
|
||||
ServerNotification::ResourceUpdatedNotification(notification) => {
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
@ -744,6 +769,10 @@ impl Peer<RoleServer> {
|
|||
}
|
||||
}
|
||||
|
||||
/// # SEP-2260: request association
|
||||
///
|
||||
/// From protocol version `2026-07-28` this must be issued while handling a
|
||||
/// client request; see [`OriginatingRequestId`].
|
||||
#[deprecated(
|
||||
since = "1.8.0",
|
||||
note = "Sampling is deprecated by SEP-2577 and will be removed in a future release. See https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2577"
|
||||
|
|
@ -778,6 +807,10 @@ impl Peer<RoleServer> {
|
|||
}
|
||||
}
|
||||
method!(
|
||||
/// # SEP-2260: request association
|
||||
///
|
||||
/// From protocol version `2026-07-28` this must be issued while handling a
|
||||
/// client request; see [`OriginatingRequestId`].
|
||||
#[deprecated(
|
||||
since = "1.8.0",
|
||||
note = "Roots is deprecated by SEP-2577 and will be removed in a future release. See https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2577"
|
||||
|
|
@ -785,9 +818,21 @@ impl Peer<RoleServer> {
|
|||
peer_req list_roots ListRootsRequest() => ListRootsResult
|
||||
);
|
||||
#[cfg(feature = "elicitation")]
|
||||
method!(peer_req create_elicitation ElicitRequest(ElicitRequestParams) => ElicitResult);
|
||||
method!(
|
||||
/// # SEP-2260: request association
|
||||
///
|
||||
/// From protocol version `2026-07-28` this must be issued while handling a
|
||||
/// client request; see [`OriginatingRequestId`].
|
||||
peer_req create_elicitation ElicitRequest(ElicitRequestParams) => ElicitResult
|
||||
);
|
||||
#[cfg(feature = "elicitation")]
|
||||
method!(peer_req_with_timeout create_elicitation_with_timeout ElicitRequest(ElicitRequestParams) => ElicitResult);
|
||||
method!(
|
||||
/// # SEP-2260: request association
|
||||
///
|
||||
/// From protocol version `2026-07-28` this must be issued while handling a
|
||||
/// client request; see [`OriginatingRequestId`].
|
||||
peer_req_with_timeout create_elicitation_with_timeout ElicitRequest(ElicitRequestParams) => ElicitResult
|
||||
);
|
||||
|
||||
method!(peer_not notify_cancelled CancelledNotification(CancelledNotificationParam));
|
||||
method!(peer_not notify_progress ProgressNotification(ProgressNotificationParam));
|
||||
|
|
@ -1014,6 +1059,11 @@ impl Peer<RoleServer> {
|
|||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
///
|
||||
/// # SEP-2260: request association
|
||||
///
|
||||
/// From protocol version `2026-07-28` this must be issued while handling a
|
||||
/// client request; see [`OriginatingRequestId`].
|
||||
#[cfg(all(feature = "schemars", feature = "elicitation"))]
|
||||
pub async fn elicit<T>(&self, message: impl Into<String>) -> Result<Option<T>, ElicitationError>
|
||||
where
|
||||
|
|
@ -1075,6 +1125,11 @@ impl Peer<RoleServer> {
|
|||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
///
|
||||
/// # SEP-2260: request association
|
||||
///
|
||||
/// From protocol version `2026-07-28` this must be issued while handling a
|
||||
/// client request; see [`OriginatingRequestId`].
|
||||
#[cfg(all(feature = "schemars", feature = "elicitation"))]
|
||||
pub async fn elicit_with_timeout<T>(
|
||||
&self,
|
||||
|
|
@ -1170,6 +1225,11 @@ impl Peer<RoleServer> {
|
|||
/// Ok(())
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// # SEP-2260: request association
|
||||
///
|
||||
/// From protocol version `2026-07-28` this must be issued while handling a
|
||||
/// client request; see [`OriginatingRequestId`].
|
||||
#[cfg(feature = "elicitation")]
|
||||
pub async fn elicit_url(
|
||||
&self,
|
||||
|
|
@ -1221,6 +1281,11 @@ impl Peer<RoleServer> {
|
|||
/// Ok(())
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// # SEP-2260: request association
|
||||
///
|
||||
/// From protocol version `2026-07-28` this must be issued while handling a
|
||||
/// client request; see [`OriginatingRequestId`].
|
||||
#[cfg(feature = "elicitation")]
|
||||
pub async fn elicit_url_with_timeout(
|
||||
&self,
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -77,6 +77,16 @@ pub enum RestoreOutcome<T> {
|
|||
/// trait for every HTTP request that carries (or should carry) a session ID.
|
||||
///
|
||||
/// See the [module-level docs](self) for background on sessions.
|
||||
///
|
||||
/// # SEP-2260 request association
|
||||
///
|
||||
/// Server-initiated requests issued while handling a client request carry an
|
||||
/// [`OriginatingRequestId`](crate::service::OriginatingRequestId) marker in
|
||||
/// their non-serialized `Extensions`; the bundled local session manager uses
|
||||
/// it to deliver such requests on the originating request's SSE stream.
|
||||
/// Implementations that serialize messages between processes lose the marker
|
||||
/// and fall back to the standalone stream, violating SEP-2260 for 2026-07-28+
|
||||
/// clients; they need their own association mechanism.
|
||||
pub trait SessionManager: Send + Sync + 'static {
|
||||
type Error: std::error::Error + Send + 'static;
|
||||
type Transport: crate::transport::Transport<RoleServer>;
|
||||
|
|
|
|||
|
|
@ -527,7 +527,36 @@ impl LocalSessionWorker {
|
|||
}
|
||||
fn resolve_outbound_channel(&self, message: &ServerJsonRpcMessage) -> OutboundChannel {
|
||||
match &message {
|
||||
ServerJsonRpcMessage::Request(_) => OutboundChannel::Common,
|
||||
// SEP-2260: requests carrying an OriginatingRequestId marker ride the
|
||||
// originating request's SSE stream, never the standalone GET stream.
|
||||
ServerJsonRpcMessage::Request(json_rpc_request) => {
|
||||
use crate::model::GetExtensions;
|
||||
let originating = json_rpc_request
|
||||
.request
|
||||
.extensions()
|
||||
.get::<crate::service::OriginatingRequestId>();
|
||||
match originating {
|
||||
Some(originating) => match self
|
||||
.resource_router
|
||||
.get(&ResourceKey::McpRequestId(originating.0.clone()))
|
||||
{
|
||||
Some(id) => OutboundChannel::RequestWise {
|
||||
id: *id,
|
||||
close: false,
|
||||
},
|
||||
None => {
|
||||
tracing::warn!(
|
||||
originating_request_id = %originating.0,
|
||||
"associated server request could not be routed to its \
|
||||
originating stream (request completed or association \
|
||||
lost); falling back to standalone stream"
|
||||
);
|
||||
OutboundChannel::Common
|
||||
}
|
||||
},
|
||||
None => OutboundChannel::Common,
|
||||
}
|
||||
}
|
||||
ServerJsonRpcMessage::Notification(JsonRpcNotification {
|
||||
notification:
|
||||
ServerNotification::ProgressNotification(Notification {
|
||||
|
|
@ -1281,3 +1310,70 @@ fn create_local_session_with_event_store(
|
|||
};
|
||||
(handle, session_worker)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod sep2260_routing_tests {
|
||||
use super::*;
|
||||
use crate::service::OriginatingRequestId;
|
||||
|
||||
fn roots_request(originating: Option<RequestId>) -> ServerJsonRpcMessage {
|
||||
#[expect(
|
||||
deprecated,
|
||||
reason = "roots is SEP-2577-deprecated; any restricted request works here"
|
||||
)]
|
||||
let mut request = crate::model::ListRootsRequest {
|
||||
method: Default::default(),
|
||||
extensions: Default::default(),
|
||||
};
|
||||
if let Some(id) = originating {
|
||||
request.extensions.insert(OriginatingRequestId(id));
|
||||
}
|
||||
ServerJsonRpcMessage::request(
|
||||
crate::model::ServerRequest::ListRootsRequest(request),
|
||||
RequestId::Number(1000),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn associated_server_request_routes_to_originating_stream() {
|
||||
let (_handle, mut worker) = create_local_session("test-session", SessionConfig::default());
|
||||
let receiver = worker.establish_request_wise_channel().await.unwrap();
|
||||
let http_request_id = receiver.http_request_id.unwrap();
|
||||
let originating_id = RequestId::Number(7);
|
||||
worker.register_resource(
|
||||
ResourceKey::McpRequestId(originating_id.clone()),
|
||||
http_request_id,
|
||||
);
|
||||
|
||||
let channel = worker.resolve_outbound_channel(&roots_request(Some(originating_id)));
|
||||
assert!(
|
||||
matches!(channel, OutboundChannel::RequestWise { id, close: false } if id == http_request_id)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unassociated_server_request_routes_to_common_stream() {
|
||||
let (_handle, mut worker) = create_local_session("test-session", SessionConfig::default());
|
||||
let _receiver = worker.establish_request_wise_channel().await.unwrap();
|
||||
let channel = worker.resolve_outbound_channel(&roots_request(None));
|
||||
assert!(matches!(channel, OutboundChannel::Common));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn associated_request_for_completed_request_falls_back_to_common() {
|
||||
// Originating request already unregistered (completion race / re-scoped
|
||||
// task op): must fall back to Common, never a closed request-wise channel.
|
||||
let (_handle, mut worker) = create_local_session("test-session", SessionConfig::default());
|
||||
let receiver = worker.establish_request_wise_channel().await.unwrap();
|
||||
let http_request_id = receiver.http_request_id.unwrap();
|
||||
let originating_id = RequestId::Number(7);
|
||||
worker.register_resource(
|
||||
ResourceKey::McpRequestId(originating_id.clone()),
|
||||
http_request_id,
|
||||
);
|
||||
worker.unregister_resource(&ResourceKey::McpRequestId(originating_id.clone()));
|
||||
|
||||
let channel = worker.resolve_outbound_channel(&roots_request(Some(originating_id)));
|
||||
assert!(matches!(channel, OutboundChannel::Common));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
273
crates/rmcp/tests/test_sep_2260_request_association.rs
Normal file
273
crates/rmcp/tests/test_sep_2260_request_association.rs
Normal file
|
|
@ -0,0 +1,273 @@
|
|||
#![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 serde_json::{Value, json};
|
||||
use tokio::{
|
||||
io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream, Lines, ReadHalf, WriteHalf},
|
||||
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(())
|
||||
}
|
||||
|
||||
// A compliant rmcp server cannot produce an unassociated 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.
|
||||
type RawServer = (
|
||||
Lines<BufReader<ReadHalf<DuplexStream>>>,
|
||||
WriteHalf<DuplexStream>,
|
||||
);
|
||||
|
||||
async fn raw_initialize(io: DuplexStream, protocol_version: &str) -> anyhow::Result<RawServer> {
|
||||
let (read, mut write) = tokio::io::split(io);
|
||||
let mut lines = BufReader::new(read).lines();
|
||||
let init: Value = serde_json::from_str(&lines.next_line().await?.expect("initialize request"))?;
|
||||
assert_eq!(init["method"], "initialize");
|
||||
let response = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": init["id"],
|
||||
"result": {
|
||||
"protocolVersion": protocol_version,
|
||||
"capabilities": {},
|
||||
"serverInfo": { "name": "raw-server", "version": "0.0.0" }
|
||||
}
|
||||
});
|
||||
write.write_all(format!("{response}\n").as_bytes()).await?;
|
||||
let initialized: Value =
|
||||
serde_json::from_str(&lines.next_line().await?.expect("initialized notification"))?;
|
||||
assert_eq!(initialized["method"], "notifications/initialized");
|
||||
Ok((lines, write))
|
||||
}
|
||||
|
||||
async fn raw_request(server: &mut RawServer, request: Value) -> anyhow::Result<Value> {
|
||||
let (lines, write) = server;
|
||||
write.write_all(format!("{request}\n").as_bytes()).await?;
|
||||
Ok(serde_json::from_str(
|
||||
&lines.next_line().await?.expect("response"),
|
||||
)?)
|
||||
}
|
||||
|
||||
fn raw_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
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unassociated_server_request_rejected_with_invalid_params() -> anyhow::Result<()> {
|
||||
let (client_io, server_io) = tokio::io::duplex(4096);
|
||||
let raw = tokio::spawn(async move {
|
||||
let mut server = raw_initialize(server_io, "2026-07-28").await?;
|
||||
raw_request(&mut server, raw_sampling_request(100)).await
|
||||
});
|
||||
|
||||
let client = SamplingClient.serve(client_io).await?;
|
||||
let response = raw.await??;
|
||||
assert_eq!(
|
||||
response["error"]["code"], -32602,
|
||||
"SEP-2260: unassociated server-to-client request must be rejected with invalid params, got {response}"
|
||||
);
|
||||
|
||||
client.cancel().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unassociated_server_request_allowed_on_legacy_protocol() -> anyhow::Result<()> {
|
||||
let (client_io, server_io) = tokio::io::duplex(4096);
|
||||
let raw = tokio::spawn(async move {
|
||||
let mut server = raw_initialize(server_io, "2025-11-25").await?;
|
||||
raw_request(&mut server, raw_sampling_request(100)).await
|
||||
});
|
||||
|
||||
let client = SamplingClient.serve(client_io).await?;
|
||||
let response = raw.await??;
|
||||
assert_eq!(
|
||||
response["result"]["model"], "test-model",
|
||||
"pre-2026-07-28 peers keep the permissive behavior, got {response}"
|
||||
);
|
||||
|
||||
client.cancel().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unassociated_ping_allowed() -> anyhow::Result<()> {
|
||||
let (client_io, server_io) = tokio::io::duplex(4096);
|
||||
let raw = tokio::spawn(async move {
|
||||
let mut server = raw_initialize(server_io, "2026-07-28").await?;
|
||||
raw_request(
|
||||
&mut server,
|
||||
json!({ "jsonrpc": "2.0", "id": 101, "method": "ping" }),
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
let client = SamplingClient.serve(client_io).await?;
|
||||
let response = raw.await??;
|
||||
assert!(
|
||||
response.get("error").is_none(),
|
||||
"SEP-2260 excepts ping from request association, got {response}"
|
||||
);
|
||||
|
||||
client.cancel().await?;
|
||||
Ok(())
|
||||
}
|
||||
165
crates/rmcp/tests/test_sep_2260_stream_routing.rs
Normal file
165
crates/rmcp/tests/test_sep_2260_stream_routing.rs
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
//! SEP-2260 end-to-end: in-handler server→client requests ride the originating
|
||||
//! POST's SSE stream, never the standalone GET stream.
|
||||
#![cfg(not(feature = "local"))]
|
||||
|
||||
use std::{collections::BTreeMap, sync::Arc, time::Duration};
|
||||
|
||||
use futures::StreamExt;
|
||||
use rmcp::{
|
||||
ErrorData as McpError, RoleServer, ServerHandler,
|
||||
model::{
|
||||
CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ElicitRequestParams,
|
||||
ElicitationSchema, ServerCapabilities, ServerInfo,
|
||||
},
|
||||
service::RequestContext,
|
||||
transport::streamable_http_server::{
|
||||
StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
|
||||
},
|
||||
};
|
||||
use serde_json::json;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct ElicitingServer;
|
||||
|
||||
impl ServerHandler for ElicitingServer {
|
||||
fn get_info(&self) -> ServerInfo {
|
||||
ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
|
||||
}
|
||||
|
||||
async fn call_tool(
|
||||
&self,
|
||||
_request: CallToolRequestParams,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> Result<CallToolResponse, McpError> {
|
||||
// Never answered: the test only checks the request is emitted on the right stream.
|
||||
let _ = context
|
||||
.peer
|
||||
.create_elicitation(ElicitRequestParams::FormElicitationParams {
|
||||
meta: None,
|
||||
message: "need input".to_string(),
|
||||
requested_schema: ElicitationSchema::new(BTreeMap::new()),
|
||||
})
|
||||
.await;
|
||||
Ok(CallToolResult::success(vec![ContentBlock::text("done")]).into())
|
||||
}
|
||||
}
|
||||
|
||||
async fn start_server(ct: CancellationToken) -> String {
|
||||
let service = StreamableHttpService::new(
|
||||
move || Ok(ElicitingServer),
|
||||
Arc::new(LocalSessionManager::default()),
|
||||
StreamableHttpServerConfig::default().with_cancellation_token(ct.child_token()),
|
||||
);
|
||||
let router = axum::Router::new().nest_service("/mcp", service);
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let url = format!(
|
||||
"http://127.0.0.1:{}/mcp",
|
||||
listener.local_addr().unwrap().port()
|
||||
);
|
||||
let ct = ct.clone();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, router)
|
||||
.with_graceful_shutdown(async move { ct.cancelled().await })
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
url
|
||||
}
|
||||
|
||||
/// Read an SSE byte stream until `needle` appears or timeout.
|
||||
async fn sse_contains(resp: reqwest::Response, needle: &str, timeout: Duration) -> bool {
|
||||
let mut stream = resp.bytes_stream();
|
||||
tokio::time::timeout(timeout, async {
|
||||
let mut buffer = String::new();
|
||||
while let Some(Ok(chunk)) = stream.next().await {
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
if buffer.contains(needle) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
})
|
||||
.await
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn elicitation_rides_originating_post_stream_not_standalone_get() {
|
||||
let ct = CancellationToken::new();
|
||||
let url = start_server(ct.clone()).await;
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
// 2025-11-25 is the latest session-carrying version; SEP-2567 serves
|
||||
// 2026-07-28+ statelessly, with no standalone GET stream to test against.
|
||||
let resp = client
|
||||
.post(&url)
|
||||
.header("Accept", "text/event-stream, application/json")
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&json!({
|
||||
"jsonrpc": "2.0", "id": 1, "method": "initialize",
|
||||
"params": {
|
||||
"protocolVersion": "2025-11-25",
|
||||
"capabilities": { "elicitation": {} },
|
||||
"clientInfo": { "name": "test-client", "version": "1.0.0" }
|
||||
}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(resp.status().is_success());
|
||||
let session_id = resp
|
||||
.headers()
|
||||
.get("Mcp-Session-Id")
|
||||
.expect("session id")
|
||||
.to_str()
|
||||
.unwrap()
|
||||
.to_string();
|
||||
|
||||
client
|
||||
.post(&url)
|
||||
.header("Accept", "text/event-stream, application/json")
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Mcp-Session-Id", &session_id)
|
||||
.json(&json!({"jsonrpc": "2.0", "method": "notifications/initialized"}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Standalone GET stream — must NEVER carry the elicitation request.
|
||||
let get_stream = client
|
||||
.get(&url)
|
||||
.header("Accept", "text/event-stream")
|
||||
.header("Mcp-Session-Id", &session_id)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(get_stream.status(), 200);
|
||||
|
||||
// The in-handler elicitation must appear on this tools/call SSE stream.
|
||||
let post_stream = client
|
||||
.post(&url)
|
||||
.header("Accept", "text/event-stream, application/json")
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Mcp-Session-Id", &session_id)
|
||||
.json(&json!({
|
||||
"jsonrpc": "2.0", "id": 2, "method": "tools/call",
|
||||
"params": { "name": "ask", "arguments": {} }
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(post_stream.status().is_success());
|
||||
|
||||
assert!(
|
||||
sse_contains(post_stream, "elicitation/create", Duration::from_secs(5)).await,
|
||||
"elicitation request must be delivered on the originating POST SSE stream"
|
||||
);
|
||||
assert!(
|
||||
!sse_contains(get_stream, "elicitation/create", Duration::from_millis(500)).await,
|
||||
"standalone GET stream must not carry the elicitation request"
|
||||
);
|
||||
|
||||
ct.cancel();
|
||||
}
|
||||
Loading…
Reference in a new issue