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 2026-09-14 14:22:57 -06:00 • committed by GitHub
commit b037c0fa88
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 726 additions and 32 deletions

View file

@ -136,6 +136,14 @@ pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone {
fn peer_cancelled_params(_notification: &Self::PeerNot) -> Option<&CancelledNotificationParam> { fn peer_cancelled_params(_notification: &Self::PeerNot) -> Option<&CancelledNotificationParam> {
None 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. /// Invalidate any response cache affected by an inbound peer notification.
/// ///
/// The serve loop calls this for every notification *before* subscription /// 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>; 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 ProgressTimeoutWatchers = Arc<tokio::sync::RwLock<HashMap<ProgressToken, mpsc::Sender<()>>>>;
type SubscriptionChannel<N> = (mpsc::Sender<N>, usize); type SubscriptionChannel<N> = (mpsc::Sender<N>, usize);
type SubscriptionChannelMap<N> = HashMap<RequestId, SubscriptionChannel<N>>; type SubscriptionChannelMap<N> = HashMap<RequestId, SubscriptionChannel<N>>;
@ -1362,7 +1374,7 @@ where
} }
let mut local_responder_pool = 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 mut local_ct_pool = HashMap::<RequestId, CancellationToken>::new();
let shared_service = Arc::new(service); let shared_service = Arc::new(service);
// for return // for return
@ -1459,7 +1471,7 @@ where
Event::SendTaskResult(SendTaskResult::Request { id, result }) => { Event::SendTaskResult(SendTaskResult::Request { id, result }) => {
if let Err(e) = result if let Err(e) = result
&& let Some(responder) = local_responder_pool.remove(&id) { && 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 { Event::SendTaskResult(SendTaskResult::Notification {
@ -1477,7 +1489,7 @@ where
&& let Some(request_id) = &param.request_id && let Some(request_id) = &param.request_id
&& let Some(responder) = local_responder_pool.remove(request_id) { && let Some(responder) = local_responder_pool.remove(request_id) {
tracing::info!(id = %request_id, reason = param.reason, "cancelled"); 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(), reason: param.reason.clone(),
})); }));
} }
@ -1514,7 +1526,10 @@ where
id, id,
responder, 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 send = transport.send(JsonRpcMessage::request(request, id.clone()));
{ {
let id = id.clone(); let id = id.clone();
@ -1615,31 +1630,38 @@ where
})) => { })) => {
tracing::info!(?notification, "received notification"); tracing::info!(?notification, "received notification");
R::invalidate_response_cache(&peer, &notification).await; R::invalidate_response_cache(&peer, &notification).await;
let cancellation_request_id = let subscription_id =
if let Some(cancelled) = R::peer_cancelled_params(&notification) { if let Some(cancelled) = R::peer_cancelled_params(&notification) {
let request_id = cancelled.request_id.clone(); let request_id = cancelled.request_id.clone();
if let Some(request_id) = request_id.as_ref() { // Modern servers cancel listen requests; legacy peers cancel
if R::IS_CLIENT { // requests they originated, even when both directions share an ID.
if let Some(responder) = if R::peer_cancels_subscriptions(&peer) {
local_responder_pool.remove(request_id) 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(), 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"); tracing::info!(id = %request_id, reason = cancelled.reason, "cancelled");
ct.cancel(); ct.cancel();
} }
None
} }
request_id
} else { } 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 if let Some(subscription_id) = subscription_id
&& let Some((sender, capacity)) = && let Some((sender, capacity)) =
peer.subscription_sender(&subscription_id) peer.subscription_sender(&subscription_id)
@ -1656,7 +1678,7 @@ where
&& let Some(responder) = && let Some(responder) =
local_responder_pool.remove(&subscription_id) local_responder_pool.remove(&subscription_id)
{ {
let _ = responder let _ = responder.responder
.send(Err(ServiceError::SubscriptionLagged { capacity })); .send(Err(ServiceError::SubscriptionLagged { capacity }));
} }
peer.unregister_subscription(&subscription_id); peer.unregister_subscription(&subscription_id);
@ -1708,7 +1730,7 @@ where
if let Some(responder) = if let Some(responder) =
remove_pending_request(&mut local_responder_pool, &id) 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 { if let Err(_error) = response_result {
tracing::warn!(%id, "Error sending response"); tracing::warn!(%id, "Error sending response");
} }
@ -1729,7 +1751,7 @@ where
} else { } else {
ServiceError::McpError(error) 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 { if let Err(_error) = _response_result {
tracing::warn!(%id, "Error sending response"); tracing::warn!(%id, "Error sending response");
} }
@ -1763,6 +1785,17 @@ where
// Then drain any handler responses still in the channel // Then drain any handler responses still in the channel
// (handlers that finished after the loop broke). // (handlers that finished after the loop broke).
while let Some(m) = sink_proxy_rx.recv().await { 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 { if let Err(error) = transport.send(m).await {
tracing::error!(%error, "failed to send pending response during drain"); tracing::error!(%error, "failed to send pending response during drain");
break; break;

View file

@ -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 // SEP-2260: reject restricted server requests that arrived unassociated
// with any in-flight outbound request. Without stream separation // with any in-flight outbound request. Without stream separation
// (`Unknown`) the coarse in-flight check under-approximates the SHOULD. // (`Unknown`) the coarse in-flight check under-approximates the SHOULD.
@ -2200,6 +2214,232 @@ where
mod tests { mod tests {
use super::*; 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] #[tokio::test]
async fn auto_startup_falls_back_when_discover_is_ignored() { async fn auto_startup_falls_back_when_discover_is_ignored() {
use crate::model::{InitializeResult, ServerCapabilities}; use crate::model::{InitializeResult, ServerCapabilities};

View file

@ -3,8 +3,22 @@
//! until the request is cancelled, so its result is only produced *after* the //! 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. //! 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}; 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::{ use rmcp::{
ErrorData as McpError, RoleServer, ServerHandler, ServiceExt, ErrorData as McpError, RoleServer, ServerHandler, ServiceExt,
model::{ model::{
@ -21,6 +35,7 @@ use tokio::{
const HELPER_ENV: &str = "RMCP_CANCELLED_RESPONSE_HELPER"; const HELPER_ENV: &str = "RMCP_CANCELLED_RESPONSE_HELPER";
const READ_TIMEOUT: Duration = Duration::from_secs(10); const READ_TIMEOUT: Duration = Duration::from_secs(10);
const STRING_REQUEST_ID: &str = "tool-request-2";
#[tokio::test(flavor = "multi_thread", worker_threads = 4)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn cancelled_request_receives_no_response() -> anyhow::Result<()> { async fn cancelled_request_receives_no_response() -> anyhow::Result<()> {
@ -43,7 +58,7 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> {
}), }),
) )
.await?; .await?;
collect_ids_until(&mut reader, 1, READ_TIMEOUT).await?; collect_ids_until(&mut reader, "1", READ_TIMEOUT).await?;
send_json( send_json(
&mut writer, &mut writer,
&json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }), &json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }),
@ -56,7 +71,7 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> {
&mut writer, &mut writer,
&json!({ &json!({
"jsonrpc": "2.0", "jsonrpc": "2.0",
"id": 2, "id": STRING_REQUEST_ID,
"method": "tools/call", "method": "tools/call",
"params": { "name": "wait-for-cancel", "arguments": {} } "params": { "name": "wait-for-cancel", "arguments": {} }
}), }),
@ -67,21 +82,59 @@ async fn cancelled_request_receives_no_response() -> anyhow::Result<()> {
&json!({ &json!({
"jsonrpc": "2.0", "jsonrpc": "2.0",
"method": "notifications/cancelled", "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?; .await?;
// A ping proves the server is alive past the cancellation, so the absence of // 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( send_json(
&mut writer, &mut writer,
&json!({ "jsonrpc": "2.0", "id": 3, "method": "ping" }), &json!({ "jsonrpc": "2.0", "id": 3, "method": "ping" }),
) )
.await?; .await?;
let seen = collect_ids_until(&mut reader, 3, READ_TIMEOUT).await?; let seen = collect_ids_until(&mut reader, "3", READ_TIMEOUT).await?;
assert!(seen.contains(&3)); assert!(seen.contains("3"));
assert!(!seen.contains(&2)); 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); drop(writer);
wait_for_child(&mut child).await; wait_for_child(&mut child).await;
@ -97,14 +150,377 @@ impl ServerHandler for WaitForCancelServer {
async fn call_tool( async fn call_tool(
&self, &self,
_request: CallToolRequestParams, request: CallToolRequestParams,
context: RequestContext<RoleServer>, context: RequestContext<RoleServer>,
) -> Result<CallToolResponse, McpError> { ) -> 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; context.ct.cancelled().await;
Ok(CallToolResult::success(vec![ContentBlock::text("late response")]).into()) 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] #[tokio::test]
async fn cancelled_response_helper() -> anyhow::Result<()> { async fn cancelled_response_helper() -> anyhow::Result<()> {
if std::env::var(HELPER_ENV).as_deref() != Ok("1") { 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. /// (then a short grace read to catch any straggler) or the timeout elapses.
async fn collect_ids_until<R>( async fn collect_ids_until<R>(
reader: &mut BufReader<R>, reader: &mut BufReader<R>,
stop_id: u64, stop_id: &str,
timeout: Duration, timeout: Duration,
) -> anyhow::Result<BTreeSet<u64>> ) -> anyhow::Result<BTreeSet<String>>
where where
R: tokio::io::AsyncRead + Unpin, R: tokio::io::AsyncRead + Unpin,
{ {
@ -200,8 +616,13 @@ where
let Ok(value) = serde_json::from_str::<Value>(trimmed) else { let Ok(value) = serde_json::from_str::<Value>(trimmed) else {
continue; continue;
}; };
if let Some(id) = value.get("id").and_then(Value::as_u64) { let id = match value.get("id") {
seen.insert(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 { if id == stop_id {
// Give any late (incorrectly-sent) response a brief window to arrive. // Give any late (incorrectly-sent) response a brief window to arrive.
deadline = tokio::time::Instant::now() + Duration::from_millis(300); deadline = tokio::time::Instant::now() + Duration::from_millis(300);