Cancel MCP tool requests when response futures are dropped (#56132)

## Why

Dropping a local tool-call future did not notify the MCP server that the request was cancelled, leaving remote work running.

## What changed

- Send a cancellation notification for the request when its response future is dropped, covering both ordinary tool calls and tool-input rounds.
- Preserve write order for stdio and in-process transports so cancellation cannot arrive before its tool request, without waiting for tool responses.

## Testing

Add regression tests for cancellation on drop, request-before-cancellation ordering under reversed send polling and exhausted task budgets, and continued use of the shared connection after cancellation.

GitOrigin-RevId: fab02d546fd82c601756117d9ecbf1017c28ac11

codex-rs/rmcp-client/src/lib.rs

@@ -23,6 +23,7 @@ mod oauth_client_credentials;
 mod oauth_client_registration;
 mod oauth_http_client;
 mod oauth_refresh_mode;
+mod ordered_byte_transport;
 mod perform_oauth_login;
 mod program_resolver;
 mod protocol_mode;
@@ -32,6 +33,7 @@ mod startup_error;
 mod stdio_server_launcher;
 mod tool_input;
 mod tool_list_changes;
+mod tool_request;
 mod trace_context;
 mod user_verification;
 mod utils;

codex-rs/rmcp-client/src/ordered_byte_transport.rs

@@ -0,0 +1,65 @@
+//! Preserve byte-stream message order even when rmcp polls send futures out of order.
+//! Only wrap transports whose send completes after writing, without awaiting a response.
+
+use std::future::Future;
+use std::io;
+use std::sync::Arc;
+
+use futures::FutureExt;
+use rmcp::RoleClient;
+use rmcp::service::RxJsonRpcMessage;
+use rmcp::service::TxJsonRpcMessage;
+use rmcp::transport::Transport;
+use tokio::sync::Semaphore;
+
+pub(crate) struct OrderedByteTransport<T> {
+    inner: T,
+    writes: Arc<Semaphore>,
+}
+
+impl<T> OrderedByteTransport<T> {
+    pub(crate) fn new(inner: T) -> Self {
+        Self {
+            inner,
+            writes: Arc::new(Semaphore::new(1)),
+        }
+    }
+}
+
+impl<T: Transport<RoleClient, Error = io::Error>> Transport<RoleClient>
+    for OrderedByteTransport<T>
+{
+    type Error = io::Error;
+
+    fn send(
+        &mut self,
+        item: TxJsonRpcMessage<RoleClient>,
+    ) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
+        // Reserve FIFO order now, before rmcp independently schedules each send.
+        // Ignore cooperative budgeting for this one poll so it always queues.
+        let mut acquire = Box::pin(Arc::clone(&self.writes).acquire_owned());
+        let acquired = tokio::task::unconstrained(acquire.as_mut()).now_or_never();
+        let send = self.inner.send(item);
+        async move {
+            let _permit = match acquired {
+                Some(result) => result,
+                None => acquire.await,
+            }
+            .map_err(io::Error::other)?;
+            // Serialize only writes, never the corresponding tool responses.
+            send.await
+        }
+    }
+
+    async fn receive(&mut self) -> Option<RxJsonRpcMessage<RoleClient>> {
+        self.inner.receive().await
+    }
+
+    async fn close(&mut self) -> Result<(), Self::Error> {
+        self.inner.close().await
+    }
+}
+
+#[cfg(test)]
+#[path = "ordered_byte_transport_tests.rs"]
+mod tests;

codex-rs/rmcp-client/src/ordered_byte_transport_tests.rs

@@ -0,0 +1,58 @@
+//! Exercise write ordering independently of rmcp's send-task scheduling.
+
+use super::OrderedByteTransport;
+use rmcp::RoleClient;
+use rmcp::RoleServer;
+use rmcp::model::CallToolRequest;
+use rmcp::model::CallToolRequestParams;
+use rmcp::model::CancelledNotification;
+use rmcp::model::CancelledNotificationParam;
+use rmcp::model::ClientJsonRpcMessage;
+use rmcp::model::ClientNotification;
+use rmcp::model::ClientRequest;
+use rmcp::model::RequestId;
+use rmcp::transport::IntoTransport;
+use rmcp::transport::Transport;
+use std::time::Duration;
+use tokio::time::timeout;
+
+#[tokio::test]
+async fn cancellation_cannot_overtake_request() -> anyhow::Result<()> {
+    let (client_io, server_io) = tokio::io::duplex(/*max_buf_size*/ 4096);
+    let mut client =
+        OrderedByteTransport::new(IntoTransport::<RoleClient, _, _>::into_transport(client_io));
+    let mut server = IntoTransport::<RoleServer, _, _>::into_transport(server_io);
+    let request = ClientJsonRpcMessage::request(
+        ClientRequest::CallToolRequest(CallToolRequest::new(CallToolRequestParams::new("wait"))),
+        RequestId::Number(1),
+    );
+    let cancellation = ClientJsonRpcMessage::notification(
+        ClientNotification::CancelledNotification(CancelledNotification::new(
+            CancelledNotificationParam::new(Some(RequestId::Number(1)), None),
+        )),
+    );
+
+    // Exhaust the task budget: reserving write order must still happen synchronously.
+    while tokio::task::coop::has_budget_remaining() {
+        tokio::task::consume_budget().await;
+    }
+    let first = client.send(request.clone());
+    let discarded = client.send(request.clone());
+    let mut second = Box::pin(client.send(cancellation.clone()));
+    tokio::task::yield_now().await;
+    assert!(futures::poll!(second.as_mut()).is_pending());
+    drop(discarded);
+    assert!(futures::poll!(second.as_mut()).is_pending());
+
+    timeout(Duration::from_secs(5), first).await??;
+    timeout(Duration::from_secs(5), second).await??;
+    let received = [
+        timeout(Duration::from_secs(5), server.receive()).await?,
+        timeout(Duration::from_secs(5), server.receive()).await?,
+    ];
+    pretty_assertions::assert_eq!(
+        serde_json::to_value(received)?,
+        serde_json::to_value([Some(request), Some(cancellation)])?
+    );
+    Ok(())
+}

codex-rs/rmcp-client/src/rmcp_client.rs

@@ -60,6 +60,7 @@ use rmcp::service::RoleClient;
 use rmcp::service::RunningService;
 use rmcp::service::ServiceError;
 use rmcp::transport::AuthorizationManager;
+use rmcp::transport::IntoTransport;
 use rmcp::transport::StreamableHttpClientTransport;
 use rmcp::transport::auth::AuthClient;
 use rmcp::transport::auth::AuthError;
@@ -94,6 +95,7 @@ use crate::oauth::validate_refresh_token_issuer;
 use crate::oauth_client_credentials::OAuthClientCredentials;
 use crate::oauth_http_client::OAuthHttpClientAdapter;
 use crate::oauth_refresh_mode::McpOAuthRefreshMode;
+use crate::ordered_byte_transport::OrderedByteTransport;
 use crate::protocol_mode::McpProtocolMode;
 use crate::startup_error::is_authentication_required_error;
 use crate::stdio_server_launcher::StdioServerCommand;
@@ -895,17 +897,9 @@ impl RmcpClient {
                     }
                     let mut options = rmcp::service::PeerRequestOptions::no_options();
                     options.meta = meta;
-                    let result = service
-                        .peer()
-                        .send_request_with_option(
-                            ClientRequest::CallToolRequest(rmcp::model::CallToolRequest::new(
-                                rmcp_params,
-                            )),
-                            options,
-                        )
-                        .await?
-                        .await_response()
-                        .await?;
+                    let (_, result) =
+                        crate::tool_request::call_tool(service.peer(), rmcp_params, options)
+                            .await?;
                     match result {
                         ServerResult::CallToolResult(result) => Ok(result),
                         _ => Err(rmcp::service::ServiceError::UnexpectedResponse),
@@ -1329,7 +1323,12 @@ impl RmcpClient {
             PendingTransport::InProcess { transport } => (
                 client_service
                     .serve_with_lifecycle(
-                        track_tool_list_changes(transport, Arc::clone(&self.tool_list_revision)),
+                        track_tool_list_changes(
+                            OrderedByteTransport::new(
+                                IntoTransport::<RoleClient, _, _>::into_transport(transport),
+                            ),
+                            Arc::clone(&self.tool_list_revision),
+                        ),
                         lifecycle,
                     )
                     .boxed(),
@@ -1338,7 +1337,10 @@ impl RmcpClient {
             PendingTransport::Stdio { transport } => (
                 client_service
                     .serve_with_lifecycle(
-                        track_tool_list_changes(*transport, Arc::clone(&self.tool_list_revision)),
+                        track_tool_list_changes(
+                            OrderedByteTransport::new(*transport),
+                            Arc::clone(&self.tool_list_revision),
+                        ),
                         lifecycle,
                     )
                     .boxed(),

codex-rs/rmcp-client/src/tool_input.rs

@@ -6,10 +6,8 @@ use std::collections::BTreeMap;
 use std::time::Duration;
 
 use rmcp::RoleClient;
-use rmcp::model::CallToolRequest;
 use rmcp::model::CallToolRequestParams;
 use rmcp::model::CallToolResult;
-use rmcp::model::ClientRequest;
 use rmcp::model::DEFAULT_MRTR_MAX_ROUNDS;
 use rmcp::model::GetExtensions;
 use rmcp::model::GetMeta;
@@ -43,17 +41,11 @@ pub(crate) async fn call_tool(
     let mut state_only_rounds = 0;
     let mut submitted_verification_proof = false;
     for round in 0..DEFAULT_MRTR_MAX_ROUNDS {
-        let response = async {
-            let handle = service
-                .peer()
-                .send_request_with_option(
-                    ClientRequest::CallToolRequest(CallToolRequest::new(params.clone())),
-                    PeerRequestOptions::no_options(),
-                )
-                .await?;
-            let id = handle.id.clone();
-            Ok::<_, ServiceError>((id, handle.await_response().await?))
-        }
+        let response = crate::tool_request::call_tool(
+            service.peer(),
+            params.clone(),
+            PeerRequestOptions::no_options(),
+        )
         .await;
         let (id, result) = response.map_err(|error| {
             if submitted_verification_proof

codex-rs/rmcp-client/src/tool_request.rs

@@ -0,0 +1,49 @@
+//! Cancel remote tool work when its local response future is dropped.
+
+use rmcp::Peer;
+use rmcp::RoleClient;
+use rmcp::model::CallToolRequest;
+use rmcp::model::CallToolRequestParams;
+use rmcp::model::CancelledNotificationParam;
+use rmcp::model::ClientRequest;
+use rmcp::model::RequestId;
+use rmcp::model::ServerResult;
+use rmcp::service::PeerRequestOptions;
+use rmcp::service::ServiceError;
+
+pub(crate) async fn call_tool(
+    peer: &Peer<RoleClient>,
+    params: CallToolRequestParams,
+    options: PeerRequestOptions,
+) -> Result<(RequestId, ServerResult), ServiceError> {
+    let handle = peer
+        .send_request_with_option(
+            ClientRequest::CallToolRequest(CallToolRequest::new(params)),
+            options,
+        )
+        .await?;
+    let id = handle.id.clone();
+    let mut cancellation = CancelOnDrop(Some((peer.clone(), id.clone())));
+    let result = handle.await_response().await;
+    cancellation.0 = None;
+    result.map(|result| (id, result))
+}
+
+struct CancelOnDrop(Option<(Peer<RoleClient>, RequestId)>);
+
+impl Drop for CancelOnDrop {
+    fn drop(&mut self) {
+        if let Some((peer, id)) = self.0.take()
+            && let Ok(runtime) = tokio::runtime::Handle::try_current()
+        {
+            runtime.spawn(async move {
+                let _ = peer
+                    .notify_cancelled(CancelledNotificationParam::new(
+                        Some(id),
+                        Some("tool call response future dropped".into()),
+                    ))
+                    .await;
+            });
+        }
+    }
+}

codex-rs/rmcp-client/src/user_verification_cancellation_tests.rs

@@ -1,4 +1,4 @@
-//! Ensures MCP cancellation releases pending elicitation routes and tool timeout pauses.
+//! Ensures MCP cancellation stops remote tool work and releases pending elicitation routes.
 
 use super::ElicitationClientService;
 use super::ElicitationPauseState;
@@ -11,21 +11,24 @@ use pretty_assertions::assert_eq;
 use rmcp::RoleServer;
 use rmcp::ServerHandler;
 use rmcp::ServiceExt;
+use rmcp::model::CallToolResult;
 use rmcp::model::CancelledNotification;
 use rmcp::model::CancelledNotificationParam;
 use rmcp::model::ClientConfig;
 use rmcp::model::ClientJsonRpcMessage;
+use rmcp::model::ClientNotification;
+use rmcp::model::ClientRequest;
 use rmcp::model::CustomRequest;
 use rmcp::model::ElicitationAction;
+use rmcp::model::InitializeResult;
 use rmcp::model::ProtocolVersion;
 use rmcp::model::RequestId;
 use rmcp::model::ServerJsonRpcMessage;
 use rmcp::model::ServerNotification;
 use rmcp::model::ServerRequest;
+use rmcp::model::ServerResult;
 use rmcp::service::PeerRequestOptions;
 use rmcp::service::QuitReason;
-use rmcp::service::RunningService;
-use rmcp::service::ServerInitializeError;
 use rmcp::service::serve_directly;
 use rmcp::transport::IntoTransport;
 use rmcp::transport::Transport;
@@ -42,9 +45,7 @@ struct ElicitationTestServer;
 impl ServerHandler for ElicitationTestServer {}
 
 struct ElicitationTestTransport {
-    servers: mpsc::UnboundedSender<
-        Result<RunningService<RoleServer, ElicitationTestServer>, ServerInitializeError>,
-    >,
+    servers: mpsc::UnboundedSender<tokio::io::DuplexStream>,
 }
 
 impl InProcessTransportFactory for ElicitationTestTransport {
@@ -52,14 +53,102 @@ impl InProcessTransportFactory for ElicitationTestTransport {
         let servers = self.servers.clone();
         Box::pin(async move {
             let (client, server) = tokio::io::duplex(/*max_buf_size*/ 4096);
-            tokio::spawn(async move {
-                let _ = servers.send(ElicitationTestServer.serve(server).await);
-            });
+            servers.send(server).map_err(std::io::Error::other)?;
             Ok(client)
         })
     }
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn dropping_tool_call_cancels_remote_request() -> anyhow::Result<()> {
+    // Run on the sole worker so independently spawned sends can be polled in reverse order.
+    tokio::spawn(async {
+        let (servers, mut server_rx) = mpsc::unbounded_channel();
+        let client = RmcpClient::new_in_process_client(Arc::new(ElicitationTestTransport { servers }))
+            .await?;
+        let mut server = IntoTransport::<RoleServer, _, _>::into_transport(server_rx.recv().await.unwrap());
+        let mut initialize = Box::pin(client.initialize(
+            ClientConfig::default().with_protocol_version(ProtocolVersion::V_2025_06_18),
+            Some(Duration::from_secs(5)),
+            Box::new(|_, _| Box::pin(async { anyhow::bail!("unexpected elicitation") })),
+        ));
+        let request = tokio::select! {
+            result = &mut initialize => panic!("client initialized before handshake: {result:?}"),
+            request = timeout(Duration::from_secs(5), server.receive()) => request?,
+        };
+        let Some(ClientJsonRpcMessage::Request(request)) = request else {
+            anyhow::bail!("expected initialize request");
+        };
+        assert!(matches!(request.request, ClientRequest::InitializeRequest(_)));
+        server
+            .send(ServerJsonRpcMessage::response(
+                ServerResult::InitializeResult(
+                    InitializeResult::default().with_protocol_version(ProtocolVersion::V_2025_06_18),
+                ),
+                request.id,
+            ))
+            .await?;
+        timeout(Duration::from_secs(5), initialize).await??;
+        let initialized = timeout(Duration::from_secs(5), server.receive()).await?;
+        assert!(matches!(
+            initialized,
+            Some(ClientJsonRpcMessage::Notification(notification))
+                if matches!(notification.notification, ClientNotification::InitializedNotification(_))
+        ));
+
+        let mut call = Box::pin(client.call_tool(
+            "wait".into(),
+            /*arguments*/ None,
+            /*meta*/ None,
+            /*timeout*/ None,
+        ));
+        // Drop before yielding to the transport's independently scheduled send tasks.
+        assert!(futures::poll!(call.as_mut()).is_pending());
+        drop(call);
+        let request = timeout(Duration::from_secs(5), server.receive()).await?;
+        let Some(ClientJsonRpcMessage::Request(request)) = request else {
+            anyhow::bail!("expected tool call before cancellation");
+        };
+        assert!(matches!(request.request, ClientRequest::CallToolRequest(_)));
+        let notification = timeout(Duration::from_secs(1), server.receive())
+            .await
+            .map_err(|_| anyhow::anyhow!("dropping the call did not cancel the server request"))?;
+        let Some(ClientJsonRpcMessage::Notification(notification)) = notification else {
+            anyhow::bail!("expected cancellation notification");
+        };
+        let ClientNotification::CancelledNotification(cancelled) = notification.notification else {
+            anyhow::bail!("expected cancellation notification");
+        };
+        assert_eq!(cancelled.params.request_id, Some(request.id));
+
+        // Cancellation applies to this request, not the shared MCP connection.
+        let mut next = Box::pin(client.call_tool(
+            "next".into(),
+            /*arguments*/ None,
+            /*meta*/ None,
+            /*timeout*/ None,
+        ));
+        let request = tokio::select! {
+            result = &mut next => panic!("tool completed before response: {result:?}"),
+            request = timeout(Duration::from_secs(5), server.receive()) => request?,
+        };
+        let Some(ClientJsonRpcMessage::Request(request)) = request else {
+            anyhow::bail!("expected next tool call");
+        };
+        let result = CallToolResult::success(Vec::new());
+        server
+            .send(ServerJsonRpcMessage::response(
+                ServerResult::CallToolResult(result.clone()),
+                request.id,
+            ))
+            .await?;
+        assert_eq!(timeout(Duration::from_secs(5), next).await??, result);
+        client.shutdown().await;
+        Ok(())
+    })
+    .await?
+}
+
 #[traced_test]
 #[tokio::test]
 async fn recovered_connections_accept_elicitations_with_previously_cancelled_ids()
@@ -77,7 +166,13 @@ async fn recovered_connections_accept_elicitations_with_previously_cancelled_ids
             "elicitationId": "authorization",
         }),
     ] {
-        let (servers, mut server_rx) = mpsc::unbounded_channel();
+        let (servers, mut transport_rx) = mpsc::unbounded_channel();
+        let (ready_tx, mut server_rx) = mpsc::unbounded_channel();
+        tokio::spawn(async move {
+            while let Some(server) = transport_rx.recv().await {
+                let _ = ready_tx.send(ElicitationTestServer.serve(server).await);
+            }
+        });
         let client =
             RmcpClient::new_in_process_client(Arc::new(ElicitationTestTransport { servers }))
                 .await?;