fix(rmcp): route peer cancellation by lifecycle (#1262)
* fix(rmcp): route peer cancellation by lifecycle Cancel inbound legacy requests by their exact request ID. Preserve modern subscription waiter cleanup using the peer lifecycle and outbound request kind. Suppress cancelled handler results and errors during EOF draining without dropping uncancelled responses during graceful shutdown. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * test(rmcp): parameterize reverse cancellation cases Use rstest to report each startup and request ID combination independently. Keep the existing per-case timeout and cancellation assertions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * fix(rmcp): use canonical initialize result in cancellation test Replace the stale ServerInfo reference after merging the upstream alias migration. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Hunter Sadler <aurokin@github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Auro
GitHub
parent
d481cb0eec
commit
b037c0fa88
3 changed files with 726 additions and 32 deletions
|
|
@ -136,6 +136,14 @@ pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone {
|
|||
fn peer_cancelled_params(_notification: &Self::PeerNot) -> Option<&CancelledNotificationParam> {
|
||||
None
|
||||
}
|
||||
#[doc(hidden)]
|
||||
fn peer_cancels_subscriptions(_peer: &Peer<Self>) -> bool {
|
||||
false
|
||||
}
|
||||
#[doc(hidden)]
|
||||
fn is_subscription_request(_request: &Self::Req) -> bool {
|
||||
false
|
||||
}
|
||||
/// Invalidate any response cache affected by an inbound peer notification.
|
||||
///
|
||||
/// The serve loop calls this for every notification *before* subscription
|
||||
|
|
@ -522,6 +530,10 @@ impl ProgressNotificationToken for ServerNotification {
|
|||
}
|
||||
|
||||
type Responder<T> = tokio::sync::oneshot::Sender<T>;
|
||||
struct PendingResponse<T> {
|
||||
responder: Responder<T>,
|
||||
is_subscription: bool,
|
||||
}
|
||||
type ProgressTimeoutWatchers = Arc<tokio::sync::RwLock<HashMap<ProgressToken, mpsc::Sender<()>>>>;
|
||||
type SubscriptionChannel<N> = (mpsc::Sender<N>, usize);
|
||||
type SubscriptionChannelMap<N> = HashMap<RequestId, SubscriptionChannel<N>>;
|
||||
|
|
@ -1362,7 +1374,7 @@ where
|
|||
}
|
||||
|
||||
let mut local_responder_pool =
|
||||
HashMap::<RequestId, Responder<Result<R::PeerResp, ServiceError>>>::new();
|
||||
HashMap::<RequestId, PendingResponse<Result<R::PeerResp, ServiceError>>>::new();
|
||||
let mut local_ct_pool = HashMap::<RequestId, CancellationToken>::new();
|
||||
let shared_service = Arc::new(service);
|
||||
// for return
|
||||
|
|
@ -1459,7 +1471,7 @@ where
|
|||
Event::SendTaskResult(SendTaskResult::Request { id, result }) => {
|
||||
if let Err(e) = result
|
||||
&& let Some(responder) = local_responder_pool.remove(&id) {
|
||||
let _ = responder.send(Err(ServiceError::TransportSend(e)));
|
||||
let _ = responder.responder.send(Err(ServiceError::TransportSend(e)));
|
||||
}
|
||||
}
|
||||
Event::SendTaskResult(SendTaskResult::Notification {
|
||||
|
|
@ -1477,7 +1489,7 @@ where
|
|||
&& let Some(request_id) = ¶m.request_id
|
||||
&& let Some(responder) = local_responder_pool.remove(request_id) {
|
||||
tracing::info!(id = %request_id, reason = param.reason, "cancelled");
|
||||
let _response_result = responder.send(Err(ServiceError::Cancelled {
|
||||
let _response_result = responder.responder.send(Err(ServiceError::Cancelled {
|
||||
reason: param.reason.clone(),
|
||||
}));
|
||||
}
|
||||
|
|
@ -1514,7 +1526,10 @@ where
|
|||
id,
|
||||
responder,
|
||||
}) => {
|
||||
local_responder_pool.insert(id.clone(), responder);
|
||||
local_responder_pool.insert(id.clone(), PendingResponse {
|
||||
responder,
|
||||
is_subscription: R::is_subscription_request(&request),
|
||||
});
|
||||
let send = transport.send(JsonRpcMessage::request(request, id.clone()));
|
||||
{
|
||||
let id = id.clone();
|
||||
|
|
@ -1615,31 +1630,38 @@ where
|
|||
})) => {
|
||||
tracing::info!(?notification, "received notification");
|
||||
R::invalidate_response_cache(&peer, ¬ification).await;
|
||||
let cancellation_request_id =
|
||||
let subscription_id =
|
||||
if let Some(cancelled) = R::peer_cancelled_params(¬ification) {
|
||||
let request_id = cancelled.request_id.clone();
|
||||
if let Some(request_id) = request_id.as_ref() {
|
||||
if R::IS_CLIENT {
|
||||
if let Some(responder) =
|
||||
local_responder_pool.remove(request_id)
|
||||
// Modern servers cancel listen requests; legacy peers cancel
|
||||
// requests they originated, even when both directions share an ID.
|
||||
if R::peer_cancels_subscriptions(&peer) {
|
||||
if let Some(request_id) = request_id.as_ref() {
|
||||
if local_responder_pool.get(request_id)
|
||||
.is_some_and(|pending| pending.is_subscription)
|
||||
&& let Some(pending) = local_responder_pool.remove(request_id)
|
||||
{
|
||||
let _ = responder.send(Err(ServiceError::Cancelled {
|
||||
let _ = pending.responder.send(Err(ServiceError::Cancelled {
|
||||
reason: cancelled.reason.clone(),
|
||||
}));
|
||||
} else {
|
||||
tracing::debug!(%request_id, "ignoring cancellation of unknown subscription");
|
||||
continue;
|
||||
}
|
||||
} else if let Some(ct) = local_ct_pool.remove(request_id) {
|
||||
}
|
||||
request_id
|
||||
} else {
|
||||
if let Some(request_id) = request_id.as_ref()
|
||||
&& let Some(ct) = local_ct_pool.remove(request_id)
|
||||
{
|
||||
tracing::info!(id = %request_id, reason = cancelled.reason, "cancelled");
|
||||
ct.cancel();
|
||||
}
|
||||
None
|
||||
}
|
||||
request_id
|
||||
} else {
|
||||
None
|
||||
notification.get_meta().subscription_id()
|
||||
};
|
||||
let subscription_id = notification
|
||||
.get_meta()
|
||||
.subscription_id()
|
||||
.or(cancellation_request_id);
|
||||
if let Some(subscription_id) = subscription_id
|
||||
&& let Some((sender, capacity)) =
|
||||
peer.subscription_sender(&subscription_id)
|
||||
|
|
@ -1656,7 +1678,7 @@ where
|
|||
&& let Some(responder) =
|
||||
local_responder_pool.remove(&subscription_id)
|
||||
{
|
||||
let _ = responder
|
||||
let _ = responder.responder
|
||||
.send(Err(ServiceError::SubscriptionLagged { capacity }));
|
||||
}
|
||||
peer.unregister_subscription(&subscription_id);
|
||||
|
|
@ -1708,7 +1730,7 @@ where
|
|||
if let Some(responder) =
|
||||
remove_pending_request(&mut local_responder_pool, &id)
|
||||
{
|
||||
let response_result = responder.send(Ok(result));
|
||||
let response_result = responder.responder.send(Ok(result));
|
||||
if let Err(_error) = response_result {
|
||||
tracing::warn!(%id, "Error sending response");
|
||||
}
|
||||
|
|
@ -1729,7 +1751,7 @@ where
|
|||
} else {
|
||||
ServiceError::McpError(error)
|
||||
};
|
||||
let _response_result = responder.send(Err(service_error));
|
||||
let _response_result = responder.responder.send(Err(service_error));
|
||||
if let Err(_error) = _response_result {
|
||||
tracing::warn!(%id, "Error sending response");
|
||||
}
|
||||
|
|
@ -1763,6 +1785,17 @@ where
|
|||
// Then drain any handler responses still in the channel
|
||||
// (handlers that finished after the loop broke).
|
||||
while let Some(m) = sink_proxy_rx.recv().await {
|
||||
if let Some(id) = match &m {
|
||||
JsonRpcMessage::Response(response) => Some(&response.id),
|
||||
JsonRpcMessage::Error(error) => error.id.as_ref(),
|
||||
_ => None,
|
||||
} {
|
||||
let Some(ct) = local_ct_pool.remove(id) else {
|
||||
tracing::debug!(%id, "dropping response for cancelled request");
|
||||
continue;
|
||||
};
|
||||
ct.cancel();
|
||||
}
|
||||
if let Err(error) = transport.send(m).await {
|
||||
tracing::error!(%error, "failed to send pending response during drain");
|
||||
break;
|
||||
|
|
|
|||
|
|
@ -292,6 +292,20 @@ impl ServiceRole for RoleClient {
|
|||
}
|
||||
}
|
||||
|
||||
fn peer_cancels_subscriptions(peer: &Peer<Self>) -> bool {
|
||||
// Discovery keeps modern lifecycle semantics even with an older application version.
|
||||
!super::uses_legacy_lifecycle(
|
||||
peer.peer_info()
|
||||
.as_deref()
|
||||
.map(|info| &info.protocol_version),
|
||||
peer.client_request_metadata.get().is_some(),
|
||||
)
|
||||
}
|
||||
|
||||
fn is_subscription_request(request: &Self::Req) -> bool {
|
||||
matches!(request, ClientRequest::SubscriptionsListenRequest(_))
|
||||
}
|
||||
|
||||
// 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.
|
||||
|
|
@ -2200,6 +2214,232 @@ where
|
|||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn server_cancellation_retires_subscription_responder() {
|
||||
tokio::task::LocalSet::new()
|
||||
.run_until(async {
|
||||
for (discover, version) in [
|
||||
(false, ProtocolVersion::V_2026_07_28),
|
||||
(true, ProtocolVersion::V_2026_07_28),
|
||||
(true, ProtocolVersion::V_2025_11_25),
|
||||
] {
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(5),
|
||||
check_subscription_cancellation(discover, version),
|
||||
)
|
||||
.await
|
||||
.expect("subscription cancellation timed out");
|
||||
}
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn check_subscription_cancellation(discover: bool, version: ProtocolVersion) {
|
||||
use crate::model::{
|
||||
GetMeta, PingRequest, ServerInfo, SubscriptionsAcknowledgedNotification,
|
||||
SubscriptionsAcknowledgedNotificationParams,
|
||||
};
|
||||
|
||||
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
||||
let mut server =
|
||||
crate::transport::IntoTransport::<RoleServer, _, _>::into_transport(server_transport);
|
||||
let info = ServerInfo {
|
||||
protocol_version: version.clone(),
|
||||
..Default::default()
|
||||
};
|
||||
let client = if discover {
|
||||
let (client, ()) = tokio::join!(
|
||||
serve_client_with_lifecycle(
|
||||
(),
|
||||
client_transport,
|
||||
ClientLifecycleMode::Discover {
|
||||
preferred_versions: vec![version.clone()]
|
||||
}
|
||||
),
|
||||
async {
|
||||
let Some(ClientJsonRpcMessage::Request(request)) = server.receive().await
|
||||
else {
|
||||
panic!("expected discover request");
|
||||
};
|
||||
assert!(matches!(request.request, ClientRequest::DiscoverRequest(_)));
|
||||
server
|
||||
.send(ServerJsonRpcMessage::response(
|
||||
ServerResult::DiscoverResult(DiscoverResult::new(
|
||||
vec![version],
|
||||
info.capabilities,
|
||||
)),
|
||||
request.id,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
);
|
||||
client.unwrap()
|
||||
} else {
|
||||
super::super::serve_directly::<RoleClient, _, _, _, _>(
|
||||
(),
|
||||
client_transport,
|
||||
Some(info.into()),
|
||||
)
|
||||
};
|
||||
let filter = SubscriptionFilter::default();
|
||||
let (subscription, ()) = tokio::join!(client.listen(filter.clone()), async {
|
||||
let Some(ClientJsonRpcMessage::Request(request)) = server.receive().await else {
|
||||
panic!("expected listen request");
|
||||
};
|
||||
assert!(matches!(
|
||||
request.request,
|
||||
ClientRequest::SubscriptionsListenRequest(_)
|
||||
));
|
||||
let mut ack: ServerNotification = SubscriptionsAcknowledgedNotification::new(
|
||||
SubscriptionsAcknowledgedNotificationParams::new(filter),
|
||||
)
|
||||
.into();
|
||||
ack.get_meta_mut().set_subscription_id(request.id);
|
||||
server
|
||||
.send(ServerJsonRpcMessage::notification(ack))
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
let mut subscription = subscription.unwrap();
|
||||
let id = subscription.id().clone();
|
||||
let mut ordinary = client
|
||||
.send_cancellable_request(
|
||||
PingRequest {
|
||||
method: Default::default(),
|
||||
extensions: Default::default(),
|
||||
}
|
||||
.into(),
|
||||
PeerRequestOptions::no_options(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
server.receive().await,
|
||||
Some(ClientJsonRpcMessage::Request(_))
|
||||
));
|
||||
for unknown in [
|
||||
ordinary.id.clone(),
|
||||
RequestId::String(id.to_string().into()),
|
||||
RequestId::String("unknown".into()),
|
||||
] {
|
||||
server
|
||||
.send(ServerJsonRpcMessage::notification(
|
||||
CancelledNotification::new(CancelledNotificationParam::new(
|
||||
Some(unknown),
|
||||
None,
|
||||
))
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
server
|
||||
.send(ServerJsonRpcMessage::request(
|
||||
PingRequest {
|
||||
method: Default::default(),
|
||||
extensions: Default::default(),
|
||||
}
|
||||
.into(),
|
||||
RequestId::String("before-cancellation".into()),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
server.receive().await,
|
||||
Some(ClientJsonRpcMessage::Response(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
subscription.request.as_mut().unwrap().rx.try_recv(),
|
||||
Err(tokio::sync::oneshot::error::TryRecvError::Empty)
|
||||
));
|
||||
assert!(matches!(
|
||||
ordinary.rx.try_recv(),
|
||||
Err(tokio::sync::oneshot::error::TryRecvError::Empty)
|
||||
));
|
||||
assert!(matches!(
|
||||
subscription.notifications.try_recv(),
|
||||
Err(tokio::sync::mpsc::error::TryRecvError::Empty)
|
||||
));
|
||||
server
|
||||
.send(ServerJsonRpcMessage::notification(
|
||||
CancelledNotification::new(CancelledNotificationParam::new(
|
||||
Some(id.clone()),
|
||||
Some("subscription ended".to_owned()),
|
||||
))
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
// Ordered input makes the ping response a barrier for cancellation handling.
|
||||
server
|
||||
.send(ServerJsonRpcMessage::request(
|
||||
PingRequest {
|
||||
method: Default::default(),
|
||||
extensions: Default::default(),
|
||||
}
|
||||
.into(),
|
||||
RequestId::String("after-cancellation".into()),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
server.receive().await,
|
||||
Some(ClientJsonRpcMessage::Response(_))
|
||||
));
|
||||
assert!(
|
||||
matches!(
|
||||
subscription.request.as_mut().unwrap().rx.try_recv(),
|
||||
Ok(Err(ServiceError::Cancelled { .. }))
|
||||
),
|
||||
"the outbound responder must retire without a final response or disconnect"
|
||||
);
|
||||
assert!(subscription.next().await.unwrap().is_none());
|
||||
assert!(matches!(
|
||||
subscription.end(),
|
||||
Some(SubscriptionEnd::Cancelled)
|
||||
));
|
||||
assert!(client.peer().subscription_sender(&id).is_none());
|
||||
server
|
||||
.send(ServerJsonRpcMessage::response(
|
||||
ServerResult::empty(()),
|
||||
ordinary.id.clone(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(ordinary.await_response().await.is_ok());
|
||||
|
||||
// Raw listen requests also retire, even without a registered notification channel.
|
||||
let raw = client
|
||||
.send_cancellable_request(
|
||||
ClientRequest::SubscriptionsListenRequest(SubscriptionsListenRequest::new(
|
||||
SubscriptionsListenRequestParams::new(SubscriptionFilter::default()),
|
||||
)),
|
||||
PeerRequestOptions::no_options(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
server.receive().await,
|
||||
Some(ClientJsonRpcMessage::Request(_))
|
||||
));
|
||||
server
|
||||
.send(ServerJsonRpcMessage::notification(
|
||||
CancelledNotification::new(CancelledNotificationParam::new(
|
||||
Some(raw.id.clone()),
|
||||
None,
|
||||
))
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
raw.await_response().await,
|
||||
Err(ServiceError::Cancelled { .. })
|
||||
));
|
||||
client.cancel().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn auto_startup_falls_back_when_discover_is_ignored() {
|
||||
use crate::model::{InitializeResult, ServerCapabilities};
|
||||
|
|
|
|||
|
|
@ -3,8 +3,22 @@
|
|||
//! until the request is cancelled, so its result is only produced *after* the
|
||||
//! cancellation — the service loop must drop it rather than write it to the wire.
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
use std::sync::Arc;
|
||||
use std::{collections::BTreeSet, process::Stdio, time::Duration};
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
use rmcp::{
|
||||
ClientHandler, RoleClient,
|
||||
model::{
|
||||
CancelledNotification, CancelledNotificationParam, ClientJsonRpcMessage, ClientRequest,
|
||||
ClientResult, ElicitRequest, ElicitRequestParams, ElicitResult, ElicitationAction,
|
||||
ElicitationSchema, PingRequest, RequestId, ServerJsonRpcMessage, ServerNotification,
|
||||
ServerRequest, ServerResult,
|
||||
},
|
||||
service::{PeerRequestOptions, QuitReason, serve_directly},
|
||||
transport::{IntoTransport, Transport},
|
||||
};
|
||||
use rmcp::{
|
||||
ErrorData as McpError, RoleServer, ServerHandler, ServiceExt,
|
||||
model::{
|
||||
|
|
@ -21,6 +35,7 @@ use tokio::{
|
|||
|
||||
const HELPER_ENV: &str = "RMCP_CANCELLED_RESPONSE_HELPER";
|
||||
const READ_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
const STRING_REQUEST_ID: &str = "tool-request-2";
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn cancelled_request_receives_no_response() -> anyhow::Result<()> {
|
||||
|
|
@ -43,7 +58,7 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> {
|
|||
}),
|
||||
)
|
||||
.await?;
|
||||
collect_ids_until(&mut reader, 1, READ_TIMEOUT).await?;
|
||||
collect_ids_until(&mut reader, "1", READ_TIMEOUT).await?;
|
||||
send_json(
|
||||
&mut writer,
|
||||
&json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }),
|
||||
|
|
@ -56,7 +71,7 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> {
|
|||
&mut writer,
|
||||
&json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 2,
|
||||
"id": STRING_REQUEST_ID,
|
||||
"method": "tools/call",
|
||||
"params": { "name": "wait-for-cancel", "arguments": {} }
|
||||
}),
|
||||
|
|
@ -67,21 +82,59 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> {
|
|||
&json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "notifications/cancelled",
|
||||
"params": { "requestId": 2 }
|
||||
"params": { "requestId": "unrelated-request" }
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
send_json(
|
||||
&mut writer,
|
||||
&json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "notifications/cancelled",
|
||||
"params": { "requestId": STRING_REQUEST_ID }
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
// A ping proves the server is alive past the cancellation, so the absence of
|
||||
// an id=2 response is genuine suppression rather than a dead connection.
|
||||
// a cancelled response is genuine suppression rather than a dead connection.
|
||||
send_json(
|
||||
&mut writer,
|
||||
&json!({ "jsonrpc": "2.0", "id": 3, "method": "ping" }),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let seen = collect_ids_until(&mut reader, 3, READ_TIMEOUT).await?;
|
||||
assert!(seen.contains(&3));
|
||||
assert!(!seen.contains(&2));
|
||||
let seen = collect_ids_until(&mut reader, "3", READ_TIMEOUT).await?;
|
||||
assert!(seen.contains("3"));
|
||||
assert!(!seen.contains(STRING_REQUEST_ID));
|
||||
|
||||
send_json(
|
||||
&mut writer,
|
||||
&json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 4,
|
||||
"method": "tools/call",
|
||||
"params": { "name": "complete-unless-cancelled", "arguments": {} }
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
send_json(
|
||||
&mut writer,
|
||||
&json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "notifications/cancelled",
|
||||
"params": { "requestId": "4" }
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
send_json(
|
||||
&mut writer,
|
||||
&json!({ "jsonrpc": "2.0", "id": 5, "method": "ping" }),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let seen = collect_ids_until(&mut reader, "5", READ_TIMEOUT).await?;
|
||||
assert!(seen.contains("4"), "string and numeric IDs must not alias");
|
||||
assert!(seen.contains("5"));
|
||||
|
||||
drop(writer);
|
||||
wait_for_child(&mut child).await;
|
||||
|
|
@ -97,14 +150,377 @@ impl ServerHandler for WaitForCancelServer {
|
|||
|
||||
async fn call_tool(
|
||||
&self,
|
||||
_request: CallToolRequestParams,
|
||||
request: CallToolRequestParams,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> Result<CallToolResponse, McpError> {
|
||||
if request.name == "complete-unless-cancelled" {
|
||||
tokio::select! {
|
||||
_ = context.ct.cancelled() => {}
|
||||
_ = tokio::time::sleep(Duration::from_millis(100)) => {}
|
||||
}
|
||||
return Ok(CallToolResult::success(vec![ContentBlock::text("completed")]).into());
|
||||
}
|
||||
context.ct.cancelled().await;
|
||||
Ok(CallToolResult::success(vec![ContentBlock::text("late response")]).into())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
#[derive(Clone)]
|
||||
struct WaitForReverseCancelClient {
|
||||
events: tokio::sync::mpsc::UnboundedSender<&'static str>,
|
||||
finish: Option<Arc<tokio::sync::Notify>>,
|
||||
return_error: bool,
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
impl ClientHandler for WaitForReverseCancelClient {
|
||||
async fn create_elicitation(
|
||||
&self,
|
||||
_request: ElicitRequestParams,
|
||||
context: rmcp::service::RequestContext<RoleClient>,
|
||||
) -> Result<ElicitResult, McpError> {
|
||||
self.events.send("started").expect("test is listening");
|
||||
context.ct.cancelled().await;
|
||||
self.events.send("cancelled").expect("test is listening");
|
||||
if let Some(finish) = &self.finish {
|
||||
finish.notified().await;
|
||||
}
|
||||
if self.return_error {
|
||||
return Err(McpError::internal_error("late handler error", None));
|
||||
}
|
||||
Ok(ElicitResult::new(ElicitationAction::Decline))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
#[rstest::rstest]
|
||||
#[case::initialize("initialize")]
|
||||
#[case::direct_legacy("direct-legacy")]
|
||||
#[case::direct_unknown("direct-unknown")]
|
||||
#[tokio::test]
|
||||
async fn peer_cancels_reverse_request_without_cancelling_outbound_request(
|
||||
#[case] startup: &str,
|
||||
#[values(false, true)] equal_ids: bool,
|
||||
) -> anyhow::Result<()> {
|
||||
tokio::time::timeout(READ_TIMEOUT, reverse_cancellation(startup, equal_ids)).await??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
async fn reverse_cancellation(startup: &str, equal_ids: bool) -> anyhow::Result<()> {
|
||||
use rmcp::model::ProtocolVersion;
|
||||
let (client_transport, server_transport) = tokio::io::duplex(4096);
|
||||
let (events_tx, mut events_rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let handler = WaitForReverseCancelClient {
|
||||
events: events_tx,
|
||||
finish: None,
|
||||
return_error: false,
|
||||
};
|
||||
let mut server = IntoTransport::<RoleServer, _, _>::into_transport(server_transport);
|
||||
let mut info = InitializeResult::default();
|
||||
info.protocol_version = ProtocolVersion::V_2025_11_25;
|
||||
let client = if startup == "initialize" {
|
||||
let (client, handshake) = tokio::join!(handler.serve(client_transport), async {
|
||||
let Some(ClientJsonRpcMessage::Request(request)) = server.receive().await else {
|
||||
panic!("expected initialize request");
|
||||
};
|
||||
assert!(matches!(
|
||||
request.request,
|
||||
ClientRequest::InitializeRequest(_)
|
||||
));
|
||||
server
|
||||
.send(ServerJsonRpcMessage::response(
|
||||
ServerResult::InitializeResult(info.clone()),
|
||||
request.id,
|
||||
))
|
||||
.await?;
|
||||
assert!(matches!(
|
||||
server.receive().await,
|
||||
Some(ClientJsonRpcMessage::Notification(_))
|
||||
));
|
||||
anyhow::Ok(())
|
||||
});
|
||||
handshake?;
|
||||
client?
|
||||
} else {
|
||||
let peer_info = (startup == "direct-legacy").then(|| info.into());
|
||||
serve_directly::<RoleClient, _, _, _, _>(handler, client_transport, peer_info)
|
||||
};
|
||||
let mut outbound = client
|
||||
.send_cancellable_request(
|
||||
ClientRequest::PingRequest(PingRequest {
|
||||
method: Default::default(),
|
||||
extensions: Default::default(),
|
||||
}),
|
||||
PeerRequestOptions::no_options(),
|
||||
)
|
||||
.await?;
|
||||
let Some(ClientJsonRpcMessage::Request(outbound_request)) = server.receive().await else {
|
||||
panic!("expected outbound ping request");
|
||||
};
|
||||
let id = if equal_ids {
|
||||
outbound_request.id.clone()
|
||||
} else {
|
||||
RequestId::String("elicitation-1".into())
|
||||
};
|
||||
let mut request = reverse_elicitation("elicitation-1");
|
||||
if let ServerJsonRpcMessage::Request(request) = &mut request {
|
||||
request.id = id.clone();
|
||||
}
|
||||
server.send(request).await?;
|
||||
assert_eq!(events_rx.recv().await, Some("started"));
|
||||
let unknown = if equal_ids {
|
||||
RequestId::String(id.to_string().into())
|
||||
} else {
|
||||
outbound_request.id.clone()
|
||||
};
|
||||
for cancel_id in [unknown, RequestId::String("unknown".into())] {
|
||||
server.send(cancel_notification(cancel_id)).await?;
|
||||
}
|
||||
exchange_ping(&mut server, "before-cancel").await?;
|
||||
assert!(matches!(
|
||||
events_rx.try_recv(),
|
||||
Err(tokio::sync::mpsc::error::TryRecvError::Empty)
|
||||
));
|
||||
assert!(matches!(
|
||||
outbound.rx.try_recv(),
|
||||
Err(tokio::sync::oneshot::error::TryRecvError::Empty)
|
||||
));
|
||||
|
||||
// A cancellation can arrive after its response was already delivered.
|
||||
server
|
||||
.send(cancel_notification(RequestId::String(
|
||||
"before-cancel".into(),
|
||||
)))
|
||||
.await?;
|
||||
server.send(cancel_notification(id)).await?;
|
||||
assert_eq!(events_rx.recv().await, Some("cancelled"));
|
||||
exchange_ping(&mut server, "after-cancel").await?;
|
||||
assert!(matches!(
|
||||
outbound.rx.try_recv(),
|
||||
Err(tokio::sync::oneshot::error::TryRecvError::Empty)
|
||||
));
|
||||
server
|
||||
.send(ServerJsonRpcMessage::response(
|
||||
ServerResult::empty(()),
|
||||
outbound_request.id,
|
||||
))
|
||||
.await?;
|
||||
assert!(matches!(
|
||||
outbound.await_response().await?,
|
||||
ServerResult::EmptyResult(_)
|
||||
));
|
||||
|
||||
client.cancel().await?;
|
||||
assert!(
|
||||
server.receive().await.is_none(),
|
||||
"no late response after handler drain"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
fn cancel_notification(id: RequestId) -> ServerJsonRpcMessage {
|
||||
ServerJsonRpcMessage::notification(
|
||||
CancelledNotification::new(CancelledNotificationParam::new(
|
||||
Some(id),
|
||||
Some("request cancelled".to_owned()),
|
||||
))
|
||||
.into(),
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
async fn exchange_ping(server: &mut impl Transport<RoleServer>, id: &str) -> anyhow::Result<()> {
|
||||
server
|
||||
.send(ServerJsonRpcMessage::request(
|
||||
PingRequest {
|
||||
method: Default::default(),
|
||||
extensions: Default::default(),
|
||||
}
|
||||
.into(),
|
||||
RequestId::String(id.into()),
|
||||
))
|
||||
.await?;
|
||||
let Some(ClientJsonRpcMessage::Response(response)) = server.receive().await else {
|
||||
panic!("expected ping response");
|
||||
};
|
||||
assert_eq!(response.id, RequestId::String(id.into()));
|
||||
assert!(matches!(response.result, ClientResult::EmptyResult(_)));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
struct ChannelClientTransport {
|
||||
incoming: tokio::sync::mpsc::UnboundedReceiver<ServerJsonRpcMessage>,
|
||||
outgoing: tokio::sync::mpsc::UnboundedSender<ClientJsonRpcMessage>,
|
||||
eof_observed: Arc<tokio::sync::Notify>,
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
impl Transport<RoleClient> for ChannelClientTransport {
|
||||
type Error = std::io::Error;
|
||||
|
||||
fn send(
|
||||
&mut self,
|
||||
item: ClientJsonRpcMessage,
|
||||
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
|
||||
let outgoing = self.outgoing.clone();
|
||||
async move {
|
||||
outgoing.send(item).map_err(|_| {
|
||||
std::io::Error::new(std::io::ErrorKind::BrokenPipe, "test receiver closed")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn receive(&mut self) -> Option<ServerJsonRpcMessage> {
|
||||
let message = self.incoming.recv().await;
|
||||
if message.is_none() {
|
||||
self.eof_observed.notify_one();
|
||||
}
|
||||
message
|
||||
}
|
||||
|
||||
async fn close(&mut self) -> Result<(), Self::Error> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
fn reverse_elicitation(id: &str) -> ServerJsonRpcMessage {
|
||||
ServerJsonRpcMessage::request(
|
||||
ServerRequest::ElicitRequest(ElicitRequest::new(
|
||||
ElicitRequestParams::FormElicitationParams {
|
||||
meta: None,
|
||||
message: "Continue the operation?".to_owned(),
|
||||
requested_schema: ElicitationSchema::builder()
|
||||
.build()
|
||||
.expect("empty elicitation schema is valid"),
|
||||
},
|
||||
)),
|
||||
RequestId::String(id.into()),
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
#[tokio::test]
|
||||
async fn graceful_close_drains_uncancelled_reverse_response() -> anyhow::Result<()> {
|
||||
const REVERSE_REQUEST_ID: &str = "elicitation-close-drain";
|
||||
|
||||
let (to_client, incoming) = tokio::sync::mpsc::unbounded_channel();
|
||||
let (outgoing, mut from_client) = tokio::sync::mpsc::unbounded_channel();
|
||||
let eof_observed = Arc::new(tokio::sync::Notify::new());
|
||||
let finish = Arc::new(tokio::sync::Notify::new());
|
||||
let (events_tx, mut events_rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let mut client = serve_directly::<RoleClient, _, _, _, _>(
|
||||
WaitForReverseCancelClient {
|
||||
events: events_tx,
|
||||
finish: Some(finish.clone()),
|
||||
return_error: false,
|
||||
},
|
||||
ChannelClientTransport {
|
||||
incoming,
|
||||
outgoing,
|
||||
eof_observed,
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
to_client.send(reverse_elicitation(REVERSE_REQUEST_ID))?;
|
||||
assert_eq!(events_rx.recv().await, Some("started"));
|
||||
|
||||
let (close_result, cancellation_event) = tokio::join!(client.close(), async {
|
||||
let event = events_rx.recv().await;
|
||||
finish.notify_one();
|
||||
event
|
||||
});
|
||||
assert!(matches!(close_result?, QuitReason::Cancelled));
|
||||
assert_eq!(cancellation_event, Some("cancelled"));
|
||||
|
||||
let Some(ClientJsonRpcMessage::Response(response)) =
|
||||
tokio::time::timeout(READ_TIMEOUT, from_client.recv()).await?
|
||||
else {
|
||||
panic!("expected in-flight response during graceful close");
|
||||
};
|
||||
assert_eq!(response.id, RequestId::String(REVERSE_REQUEST_ID.into()));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
#[tokio::test]
|
||||
async fn cancelled_reverse_request_stays_suppressed_during_eof_drain() -> anyhow::Result<()> {
|
||||
for return_error in [false, true] {
|
||||
tokio::time::timeout(
|
||||
READ_TIMEOUT,
|
||||
cancelled_reverse_response_at_eof(return_error),
|
||||
)
|
||||
.await??;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "client", not(feature = "local")))]
|
||||
async fn cancelled_reverse_response_at_eof(return_error: bool) -> anyhow::Result<()> {
|
||||
const REVERSE_REQUEST_ID: &str = "elicitation-before-eof";
|
||||
|
||||
let (to_client, incoming) = tokio::sync::mpsc::unbounded_channel();
|
||||
let (outgoing, mut from_client) = tokio::sync::mpsc::unbounded_channel();
|
||||
let eof_observed = Arc::new(tokio::sync::Notify::new());
|
||||
let finish = Arc::new(tokio::sync::Notify::new());
|
||||
let (events_tx, mut events_rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let client = serve_directly::<RoleClient, _, _, _, _>(
|
||||
WaitForReverseCancelClient {
|
||||
events: events_tx,
|
||||
finish: Some(finish.clone()),
|
||||
return_error,
|
||||
},
|
||||
ChannelClientTransport {
|
||||
incoming,
|
||||
outgoing,
|
||||
eof_observed: eof_observed.clone(),
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
to_client.send(reverse_elicitation(REVERSE_REQUEST_ID))?;
|
||||
assert_eq!(
|
||||
tokio::time::timeout(READ_TIMEOUT, events_rx.recv()).await?,
|
||||
Some("started")
|
||||
);
|
||||
|
||||
to_client.send(ServerJsonRpcMessage::notification(
|
||||
ServerNotification::CancelledNotification(CancelledNotification::new(
|
||||
CancelledNotificationParam::new(
|
||||
Some(RequestId::String(REVERSE_REQUEST_ID.into())),
|
||||
Some("user cancelled".to_owned()),
|
||||
),
|
||||
)),
|
||||
))?;
|
||||
assert_eq!(
|
||||
tokio::time::timeout(READ_TIMEOUT, events_rx.recv()).await?,
|
||||
Some("cancelled")
|
||||
);
|
||||
|
||||
let eof = eof_observed.notified();
|
||||
drop(to_client);
|
||||
tokio::time::timeout(READ_TIMEOUT, eof).await?;
|
||||
finish.notify_one();
|
||||
|
||||
assert!(
|
||||
matches!(
|
||||
tokio::time::timeout(READ_TIMEOUT, client.waiting()).await??,
|
||||
QuitReason::Closed
|
||||
),
|
||||
"client must close after input EOF"
|
||||
);
|
||||
assert!(
|
||||
from_client.recv().await.is_none(),
|
||||
"cancelled reverse request must not be sent during EOF drain (return_error={return_error})"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancelled_response_helper() -> anyhow::Result<()> {
|
||||
if std::env::var(HELPER_ENV).as_deref() != Ok("1") {
|
||||
|
|
@ -172,9 +588,9 @@ where
|
|||
/// (then a short grace read to catch any straggler) or the timeout elapses.
|
||||
async fn collect_ids_until<R>(
|
||||
reader: &mut BufReader<R>,
|
||||
stop_id: u64,
|
||||
stop_id: &str,
|
||||
timeout: Duration,
|
||||
) -> anyhow::Result<BTreeSet<u64>>
|
||||
) -> anyhow::Result<BTreeSet<String>>
|
||||
where
|
||||
R: tokio::io::AsyncRead + Unpin,
|
||||
{
|
||||
|
|
@ -200,8 +616,13 @@ where
|
|||
let Ok(value) = serde_json::from_str::<Value>(trimmed) else {
|
||||
continue;
|
||||
};
|
||||
if let Some(id) = value.get("id").and_then(Value::as_u64) {
|
||||
seen.insert(id);
|
||||
let id = match value.get("id") {
|
||||
Some(Value::String(id)) => Some(id.clone()),
|
||||
Some(Value::Number(id)) => Some(id.to_string()),
|
||||
_ => None,
|
||||
};
|
||||
if let Some(id) = id {
|
||||
seen.insert(id.clone());
|
||||
if id == stop_id {
|
||||
// Give any late (incorrectly-sent) response a brief window to arrive.
|
||||
deadline = tokio::time::Instant::now() + Duration::from_millis(300);
|
||||
|
|
|
|||
Loading…
Reference in a new issue