diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 772844d86..f0590a6ca 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -371,6 +371,11 @@ name = "test_trace_context" required-features = ["server", "client"] path = "tests/test_trace_context.rs" +[[test]] +name = "test_abandoned_request_cancel" +required-features = ["server", "client"] +path = "tests/test_abandoned_request_cancel.rs" + [[test]] name = "test_mrtr_behavior" required-features = ["server", "client"] diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 5dad6554a..fd8c2d286 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -448,6 +448,39 @@ impl Drop for PendingSubscriptionRequest { } } +/// Tells the server about a request whose caller stopped waiting for it. +/// +/// Armed while the response is awaited. Dropping it armed, which happens when +/// the caller drops the future (a `tokio::time::timeout`, a lost `select!` +/// branch), sends a best-effort `notifications/cancelled` for the request, so +/// the server does not keep working on a result nobody will read. +struct AbandonedRequestGuard<'a> { + peer: &'a Peer, + id: Option, +} + +impl<'a> AbandonedRequestGuard<'a> { + const REASON: &'static str = "request abandoned by the caller"; + + fn new(peer: &'a Peer, id: RequestId) -> Self { + Self { peer, id: Some(id) } + } + + fn disarm(&mut self) { + self.id.take(); + } +} + +impl Drop for AbandonedRequestGuard<'_> { + fn drop(&mut self) { + let Some(id) = self.id.take() else { + return; + }; + self.peer + .try_cancel_request(id, Some(Self::REASON.to_owned())); + } +} + impl Subscription { /// Return the originating listen request ID. pub fn id(&self) -> &RequestId { @@ -1110,7 +1143,7 @@ macro_rules! method { $(#[$meta])* pub async fn $method(&self) -> Result<$Resp, ServiceError> { let result = self - .send_request(ClientRequest::$Req($Req { + .send_request_cancel_on_drop(ClientRequest::$Req($Req { method: Default::default(), })) .await?; @@ -1124,7 +1157,7 @@ macro_rules! method { $(#[$meta])* pub async fn $method(&self, params: $Param) -> Result<$Resp, ServiceError> { let result = self - .send_request(ClientRequest::$Req($Req { + .send_request_cancel_on_drop(ClientRequest::$Req($Req { method: Default::default(), params, extensions: Default::default(), @@ -1140,7 +1173,7 @@ macro_rules! method { $(#[$meta])* pub async fn $method(&self, params: Option<$Param>) -> Result<$Resp, ServiceError> { let result = self - .send_request(ClientRequest::$Req($Req { + .send_request_cancel_on_drop(ClientRequest::$Req($Req { method: Default::default(), params, extensions: Default::default(), @@ -1156,7 +1189,7 @@ macro_rules! method { $(#[$meta])* pub async fn $method(&self, params: $Param) -> Result<(), ServiceError> { let result = self - .send_request(ClientRequest::$Req($Req { + .send_request_cancel_on_drop(ClientRequest::$Req($Req { method: Default::default(), params, extensions: Default::default(), @@ -1374,6 +1407,21 @@ impl Peer { .await; } + /// [`Peer::send_request`] that sends `notifications/cancelled` if dropped before + /// the response arrives. Never route `initialize` or `server/discover` through here. + async fn send_request_cancel_on_drop( + &self, + request: ClientRequest, + ) -> Result { + let handle = self + .send_request_with_option(request, PeerRequestOptions::no_options()) + .await?; + let mut guard = AbandonedRequestGuard::new(self, handle.id.clone()); + let result = handle.await_response().await; + guard.disarm(); + result + } + /// Send one `tools/call` request and return either a final result or an MRTR /// `InputRequiredResult` without driving any follow-up rounds. pub async fn call_tool_once( @@ -1381,7 +1429,7 @@ impl Peer { params: CallToolRequestParams, ) -> Result { let result = self - .send_request(ClientRequest::CallToolRequest(CallToolRequest { + .send_request_cancel_on_drop(ClientRequest::CallToolRequest(CallToolRequest { method: Default::default(), params, extensions: Default::default(), @@ -1401,7 +1449,7 @@ impl Peer { /// SEP-2663 `tasks/get`: poll the current state of a task. pub async fn get_task(&self, params: GetTaskParams) -> Result { let result = self - .send_request(ClientRequest::GetTaskRequest(GetTaskRequest::new(params))) + .send_request_cancel_on_drop(ClientRequest::GetTaskRequest(GetTaskRequest::new(params))) .await?; match result { ServerResult::GetTaskResult(result) => Ok(result), @@ -1413,7 +1461,7 @@ impl Peer { /// input requests. The acknowledgement is eventually consistent. pub async fn update_task(&self, params: UpdateTaskParams) -> Result<(), ServiceError> { let result = self - .send_request(ClientRequest::UpdateTaskRequest(UpdateTaskRequest::new( + .send_request_cancel_on_drop(ClientRequest::UpdateTaskRequest(UpdateTaskRequest::new( params, ))) .await?; @@ -1427,7 +1475,7 @@ impl Peer { /// is cooperative; the ack does not guarantee the task stops. pub async fn cancel_task(&self, params: CancelTaskParams) -> Result<(), ServiceError> { let result = self - .send_request(ClientRequest::CancelTaskRequest(CancelTaskRequest::new( + .send_request_cancel_on_drop(ClientRequest::CancelTaskRequest(CancelTaskRequest::new( params, ))) .await?; @@ -1444,7 +1492,7 @@ impl Peer { params: GetPromptRequestParams, ) -> Result { let result = self - .send_request(ClientRequest::GetPromptRequest(GetPromptRequest { + .send_request_cancel_on_drop(ClientRequest::GetPromptRequest(GetPromptRequest { method: Default::default(), params, extensions: Default::default(), @@ -1474,7 +1522,7 @@ impl Peer { let generation = self.capture_response_cache_generation().await; let result = self - .send_request(ClientRequest::ReadResourceRequest(ReadResourceRequest { + .send_request_cancel_on_drop(ClientRequest::ReadResourceRequest(ReadResourceRequest { method: Default::default(), params, extensions: Default::default(), @@ -1547,7 +1595,7 @@ impl Peer { let generation = self.capture_response_cache_generation().await; let uses_cursor = request_uses_cursor(¶ms); let result = self - .send_request(ClientRequest::ListPromptsRequest(ListPromptsRequest { + .send_request_cancel_on_drop(ClientRequest::ListPromptsRequest(ListPromptsRequest { method: Default::default(), params, extensions: Default::default(), @@ -1597,11 +1645,13 @@ impl Peer { let generation = self.capture_response_cache_generation().await; let uses_cursor = request_uses_cursor(¶ms); let result = self - .send_request(ClientRequest::ListResourcesRequest(ListResourcesRequest { - method: Default::default(), - params, - extensions: Default::default(), - })) + .send_request_cancel_on_drop(ClientRequest::ListResourcesRequest( + ListResourcesRequest { + method: Default::default(), + params, + extensions: Default::default(), + }, + )) .await; let result = match result { Ok(result) => result, @@ -1648,7 +1698,7 @@ impl Peer { let generation = self.capture_response_cache_generation().await; let uses_cursor = request_uses_cursor(¶ms); let result = self - .send_request(ClientRequest::ListResourceTemplatesRequest( + .send_request_cancel_on_drop(ClientRequest::ListResourceTemplatesRequest( ListResourceTemplatesRequest { method: Default::default(), params, @@ -1710,7 +1760,7 @@ impl Peer { let generation = self.capture_response_cache_generation().await; let uses_cursor = request_uses_cursor(¶ms); let result = self - .send_request(ClientRequest::ListToolsRequest(ListToolsRequest { + .send_request_cancel_on_drop(ClientRequest::ListToolsRequest(ListToolsRequest { method: Default::default(), params, extensions: Default::default(), diff --git a/crates/rmcp/tests/test_abandoned_request_cancel.rs b/crates/rmcp/tests/test_abandoned_request_cancel.rs new file mode 100644 index 000000000..0bd272084 --- /dev/null +++ b/crates/rmcp/tests/test_abandoned_request_cancel.rs @@ -0,0 +1,386 @@ +//! A client that stops waiting for a request tells the server with +//! `notifications/cancelled`. +//! +//! The request helpers hide the request they send, and the MRTR helpers +//! (`call_tool`, `get_prompt`, `read_resource`) hide one per round, so the only +//! way a caller can abandon one is to drop the future. These tests drop it at +//! each point of a round and check what the server is told. + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use rmcp::{ + ClientHandler, ServerHandler, + model::*, + service::{ + NotificationContext, RequestContext, RoleClient, RoleServer, RunningService, serve_directly, + }, +}; +use serde_json::json; +use tokio::sync::{Notify, mpsc}; + +const ABANDONED_REASON: &str = "request abandoned by the caller"; +const WAIT: Duration = Duration::from_secs(5); + +/// Tool, prompt and resource names select the behavior: +/// - `hang`: blocks until the request is cancelled; +/// - `ask_then_hang`: one elicitation round, then blocks on the retry; +/// - `ask_blocking`: one elicitation the client never answers; +/// - anything else completes at once. +/// +/// `prompts/list` and `completion/complete` always block, to stand for the +/// requests that have no MRTR rounds. +#[derive(Clone)] +struct ProbeServer { + /// Signalled when a request starts to block. + hanging: Arc, + /// ID of the request that is blocking. + hanging_id: Arc>>, + /// Every `notifications/cancelled` the server receives. + cancelled: mpsc::UnboundedSender, +} + +impl ProbeServer { + async fn hang(&self, context: RequestContext) -> ErrorData { + *self.hanging_id.lock().unwrap() = Some(context.id.clone()); + self.hanging.notify_one(); + context.ct.cancelled().await; + ErrorData::internal_error("cancelled", None) + } +} + +fn elicitation(message: &str) -> InputRequiredResult { + let mut requests = InputRequests::new(); + requests.insert( + "answer".to_string(), + InputRequest::Elicitation(ElicitRequest::new( + ElicitRequestParams::FormElicitationParams { + meta: None, + message: message.into(), + requested_schema: serde_json::from_value(json!({ + "type": "object", + "properties": { "name": { "type": "string" } }, + "required": ["name"] + })) + .unwrap(), + }, + )), + ); + InputRequiredResult::new(Some(requests), Some("state".into())) +} + +impl ServerHandler for ProbeServer { + fn get_info(&self) -> ServerConfig { + let mut info = ServerConfig::new( + ServerCapabilities::builder() + .enable_tools() + .enable_prompts() + .enable_resources() + .build(), + ); + info.protocol_version = ProtocolVersion::V_2026_07_28; + info + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + let answered = request.input_responses.is_some(); + match request.name.as_ref() { + "hang" => Err(self.hang(context).await), + "ask_then_hang" if answered => Err(self.hang(context).await), + "ask_then_hang" => Ok(elicitation("Name?").into()), + "ask_blocking" => Ok(elicitation("BLOCK").into()), + _ => Ok(CallToolResult::success(vec![ContentBlock::text("done")]).into()), + } + } + + async fn get_prompt( + &self, + _request: GetPromptRequestParams, + context: RequestContext, + ) -> Result { + Err(self.hang(context).await) + } + + async fn read_resource( + &self, + _request: ReadResourceRequestParams, + context: RequestContext, + ) -> Result { + Err(self.hang(context).await) + } + + async fn list_prompts( + &self, + _request: Option, + context: RequestContext, + ) -> Result { + Err(self.hang(context).await) + } + + async fn complete( + &self, + _request: CompleteRequestParams, + context: RequestContext, + ) -> Result { + Err(self.hang(context).await) + } + + async fn on_cancelled( + &self, + notification: CancelledNotificationParam, + _context: NotificationContext, + ) { + let _ = self.cancelled.send(notification); + } +} + +/// Answers every elicitation except `"BLOCK"`, which it never answers. +#[derive(Clone)] +struct ProbeClient { + /// Signalled when the blocking elicitation starts. + eliciting: Arc, +} + +impl ClientHandler for ProbeClient { + async fn create_elicitation( + &self, + request: ElicitRequestParams, + _context: RequestContext, + ) -> Result { + if let ElicitRequestParams::FormElicitationParams { message, .. } = &request + && message == "BLOCK" + { + self.eliciting.notify_one(); + std::future::pending::<()>().await; + } + Ok(ElicitResult::new(ElicitationAction::Accept).with_content(json!({ "name": "Ferris" }))) + } +} + +struct Probe { + client: RunningService, + hanging: Arc, + hanging_id: Arc>>, + eliciting: Arc, + cancelled: mpsc::UnboundedReceiver, +} + +impl Probe { + /// The single cancellation the server should have received. + async fn expect_cancelled(&mut self) -> CancelledNotificationParam { + tokio::time::timeout(WAIT, self.cancelled.recv()) + .await + .expect("server received no notifications/cancelled") + .expect("server dropped the cancellation channel") + } + + /// Runs a request to completion after the point under test, so that any + /// cancellation sent before it has reached the server, then checks that + /// none did. + async fn expect_no_cancellation(&mut self) -> anyhow::Result<()> { + self.client + .call_tool(CallToolRequestParams::new("noop")) + .await?; + tokio::time::sleep(Duration::from_millis(50)).await; + assert!( + self.cancelled.try_recv().is_err(), + "server received an unexpected notifications/cancelled" + ); + Ok(()) + } +} + +/// Runs `body` inside a `LocalSet` so `spawn_local` (used when the `local` +/// feature is active) is available, wiring up a connected client/server pair. +async fn with_probe(body: F) -> anyhow::Result<()> +where + F: FnOnce(Probe) -> Fut, + Fut: std::future::Future>, +{ + tokio::task::LocalSet::new() + .run_until(async move { + let (cancelled_tx, cancelled) = mpsc::unbounded_channel(); + let server = ProbeServer { + hanging: Arc::new(Notify::new()), + hanging_id: Arc::new(Mutex::new(None)), + cancelled: cancelled_tx, + }; + let hanging = server.hanging.clone(); + let hanging_id = server.hanging_id.clone(); + let eliciting = Arc::new(Notify::new()); + + let (server_transport, client_transport) = tokio::io::duplex(8192); + let server_peer_info = InitializeRequestParams::new( + ClientCapabilities::builder().enable_elicitation().build(), + Implementation::new("abandoned-request-test-client", "0.0.0"), + ) + .with_protocol_version(ProtocolVersion::V_2026_07_28); + let mut client_peer_info = + InitializeResult::new(ServerCapabilities::builder().enable_tools().build()); + client_peer_info.protocol_version = ProtocolVersion::V_2026_07_28; + + let server_task = tokio::task::spawn_local(async move { + let running = serve_directly::( + server, + server_transport, + Some(server_peer_info), + ); + running.waiting().await?; + anyhow::Ok(()) + }); + let client = serve_directly::( + ProbeClient { + eliciting: eliciting.clone(), + }, + client_transport, + Some(client_peer_info.into()), + ); + + let result = body(Probe { + client, + hanging, + hanging_id, + eliciting, + cancelled, + }) + .await; + + server_task.abort(); + result + }) + .await +} + +/// Polls `call` until `reached` fires, then drops it: the caller gave up at +/// that point. +async fn abandon_when(call: impl Future, reached: &Notify) { + let mut call = Box::pin(call); + tokio::select! { + _ = &mut call => panic!("the call completed before it could be abandoned"), + reached = tokio::time::timeout(WAIT, reached.notified()) => { + reached.expect("the call never reached the point to abandon it at"); + } + } + drop(call); +} + +#[tokio::test(flavor = "current_thread")] +async fn abandoned_tool_call_is_cancelled() -> anyhow::Result<()> { + with_probe(|mut probe| async move { + abandon_when( + probe.client.call_tool(CallToolRequestParams::new("hang")), + &probe.hanging, + ) + .await; + + let cancelled = probe.expect_cancelled().await; + assert_eq!(cancelled.request_id, *probe.hanging_id.lock().unwrap()); + assert_eq!(cancelled.reason.as_deref(), Some(ABANDONED_REASON)); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn abandoned_retry_round_cancels_that_round() -> anyhow::Result<()> { + with_probe(|mut probe| async move { + // The first round completes with `input_required`; the call is dropped + // while the retry is in flight, so the retry is the request to cancel. + abandon_when( + probe + .client + .call_tool(CallToolRequestParams::new("ask_then_hang")), + &probe.hanging, + ) + .await; + + let cancelled = probe.expect_cancelled().await; + assert_eq!(cancelled.request_id, *probe.hanging_id.lock().unwrap()); + probe.expect_no_cancellation().await + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn abandoned_prompt_and_resource_requests_are_cancelled() -> anyhow::Result<()> { + with_probe(|mut probe| async move { + abandon_when( + probe.client.get_prompt(GetPromptRequestParams::new("hang")), + &probe.hanging, + ) + .await; + let cancelled = probe.expect_cancelled().await; + assert_eq!(cancelled.request_id, *probe.hanging_id.lock().unwrap()); + + abandon_when( + probe + .client + .read_resource(ReadResourceRequestParams::new("test://hang")), + &probe.hanging, + ) + .await; + let cancelled = probe.expect_cancelled().await; + assert_eq!(cancelled.request_id, *probe.hanging_id.lock().unwrap()); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn abandoned_requests_without_mrtr_rounds_are_cancelled() -> anyhow::Result<()> { + with_probe(|mut probe| async move { + abandon_when(probe.client.list_prompts(None), &probe.hanging).await; + let cancelled = probe.expect_cancelled().await; + assert_eq!(cancelled.request_id, *probe.hanging_id.lock().unwrap()); + assert_eq!(cancelled.reason.as_deref(), Some(ABANDONED_REASON)); + + let complete = CompleteRequestParams::new( + Reference::for_prompt("anything"), + ArgumentInfo::new("argument", "va"), + ); + abandon_when(probe.client.complete(complete), &probe.hanging).await; + let cancelled = probe.expect_cancelled().await; + assert_eq!(cancelled.request_id, *probe.hanging_id.lock().unwrap()); + probe.expect_no_cancellation().await + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn abandoning_during_local_input_handling_cancels_nothing() -> anyhow::Result<()> { + with_probe(|mut probe| async move { + // The round's request has completed with `input_required` and the retry + // has not been sent: no request is in flight while the client handler + // works on the elicitation, so there is nothing to cancel. + abandon_when( + probe + .client + .call_tool(CallToolRequestParams::new("ask_blocking")), + &probe.eliciting, + ) + .await; + + probe.expect_no_cancellation().await + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn completed_call_cancels_nothing() -> anyhow::Result<()> { + with_probe(|mut probe| async move { + let result = probe + .client + .call_tool(CallToolRequestParams::new("noop")) + .await?; + assert_eq!(result.content[0].as_text().unwrap().text, "done"); + + probe.expect_no_cancellation().await + }) + .await +}