Actually send cancellation to MCP servers (#3865)

crates/mcp-client/src/client.rs

@@ -1,19 +1,22 @@
 use rmcp::{
     model::{
-        CallToolRequest, CallToolRequestParam, CallToolResult, ClientCapabilities, ClientInfo,
+        CallToolRequest, CallToolRequestParam, CallToolResult, CancelledNotification,
+        CancelledNotificationMethod, CancelledNotificationParam, ClientCapabilities, ClientInfo,
         ClientRequest, GetPromptRequest, GetPromptRequestParam, GetPromptResult, Implementation,
         InitializeResult, ListPromptsRequest, ListPromptsResult, ListResourcesRequest,
         ListResourcesResult, ListToolsRequest, ListToolsResult, LoggingMessageNotification,
         LoggingMessageNotificationMethod, PaginatedRequestParam, ProgressNotification,
         ProgressNotificationMethod, ProtocolVersion, ReadResourceRequest, ReadResourceRequestParam,
-        ReadResourceResult, ServerNotification, ServerResult,
+        ReadResourceResult, RequestId, ServerNotification, ServerResult,
+    },
+    service::{
+        ClientInitializeError, PeerRequestOptions, RequestHandle, RunningService, ServiceRole,
     },
-    service::{ClientInitializeError, PeerRequestOptions, RunningService},
     transport::IntoTransport,
-    ClientHandler, RoleClient, ServiceError, ServiceExt,
+    ClientHandler, Peer, RoleClient, ServiceError, ServiceExt,
 };
 use serde_json::Value;
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 use tokio::sync::{
     mpsc::{self, Sender},
     Mutex,
@@ -176,27 +179,52 @@ impl McpClient {
             .client
             .lock()
             .await
-            .send_request_with_option(
+            .send_cancellable_request(request, PeerRequestOptions::no_options())
-                request,
-                PeerRequestOptions {
-                    timeout: Some(self.timeout),
-                    meta: None,
-                },
-            )
             .await?;
 
-        let cancel_token = cancel_token.clone();
+        await_response(handle, self.timeout, &cancel_token).await
-        tokio::select! {
+    }
-            res = handle.await_response() => {
+}
-                Ok(res?)
+
-            }
+async fn await_response(
-            _ = cancel_token.cancelled() => {
+    handle: RequestHandle<RoleClient>,
-                Err(Error::Cancelled{reason: None})
+    timeout: Duration,
-            }
+    cancel_token: &CancellationToken,
+) -> Result<<RoleClient as ServiceRole>::PeerResp, ServiceError> {
+    let receiver = handle.rx;
+    let peer = handle.peer;
+    let request_id = handle.id;
+    tokio::select! {
+        result = receiver => {
+            result.map_err(|_e| ServiceError::TransportClosed)?
+        }
+        _ = tokio::time::sleep(timeout) => {
+            send_cancel_message(&peer, request_id, Some("timed out".to_owned())).await?;
+            Err(ServiceError::Timeout{timeout})
+        }
+        _ = cancel_token.cancelled() => {
+            send_cancel_message(&peer, request_id, Some("operation cancelled".to_owned())).await?;
+            Err(ServiceError::Cancelled { reason: None })
         }
     }
 }
 
+async fn send_cancel_message(
+    peer: &Peer<RoleClient>,
+    request_id: RequestId,
+    reason: Option<String>,
+) -> Result<(), ServiceError> {
+    peer.send_notification(
+        CancelledNotification {
+            params: CancelledNotificationParam { request_id, reason },
+            method: CancelledNotificationMethod,
+            extensions: Default::default(),
+        }
+        .into(),
+    )
+    .await
+}
+
 #[async_trait::async_trait]
 impl McpClientTrait for McpClient {
     fn get_info(&self) -> Option<&InitializeResult> {