From b427e9771a24e9367cf069340631cd9451f97905 Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 15 Oct 2026 11:26:35 -0400 Subject: [PATCH] fix(transport): preserve protocol error replies across cancellation --- crates/rmcp/src/transport/async_rw.rs | 115 ++++++++++++++++++++++++-- 1 file changed, 106 insertions(+), 9 deletions(-) diff --git a/crates/rmcp/src/transport/async_rw.rs b/crates/rmcp/src/transport/async_rw.rs index 59d5eccdd..240c3aa55 100644 --- a/crates/rmcp/src/transport/async_rw.rs +++ b/crates/rmcp/src/transport/async_rw.rs @@ -66,9 +66,15 @@ pub struct AsyncRwTransport { /// `receive` is polled inside a `select!` and can be cancelled mid-discard. discarding: bool, write: Arc>>>, + pending_error_response: Option, _role: PhantomData Role>, } +enum PendingErrorResponse { + Queue, + Flush, +} + /// Outcome of a single bounded line read. enum LineRead { /// A whole `\n`-terminated line is in `line_buf`. @@ -97,6 +103,7 @@ where max_line_length: DEFAULT_MAX_LINE_LENGTH, discarding: false, write, + pending_error_response: None, _role: PhantomData, } } @@ -244,6 +251,26 @@ where async fn receive(&mut self) -> Option> { loop { + // Keep reply progress across cancellation at either the writer lock or + // flush. Once queued, resuming must flush rather than enqueue a duplicate. + if let Some(state) = self.pending_error_response.as_mut() { + let mut write = self.write.lock().await; + let framed = write.as_mut()?; + if matches!(state, PendingErrorResponse::Queue) { + let response = TxJsonRpcMessage::::error( + ErrorData::invalid_request("Invalid request", None), + None, + ); + if framed.feed(response).await.is_err() { + return None; + } + *state = PendingErrorResponse::Flush; + } + if framed.flush().await.is_err() { + return None; + } + self.pending_error_response = None; + } // `receive` is polled inside a `select!` in the service loop, so an // in-progress line read is dropped whenever another branch (e.g. an // outgoing response) becomes ready. `read_line` appends into @@ -302,15 +329,7 @@ where // real protocol error rather than unparsable input, so surface it with // an Invalid Request response instead of silently dropping it. tracing::debug!("Protocol error on incoming message: {e}"); - let mut write = self.write.lock().await; - let framed = write.as_mut()?; - let response = TxJsonRpcMessage::::error( - ErrorData::invalid_request("Invalid request", None), - None, - ); - if framed.send(response).await.is_err() { - return None; - } + self.pending_error_response = Some(PendingErrorResponse::Queue); } } } @@ -902,4 +921,82 @@ mod test { }), ); } + + #[cfg(feature = "server")] + #[tokio::test] + async fn protocol_error_reply_survives_cancellation_waiting_for_writer() { + use tokio::io::AsyncReadExt; + + use crate::RoleServer; + + let input = b"{\"foo\":\"bar\"}\n{\"jsonrpc\":\"2.0\",\"method\":\"ping\",\"id\":1}\n"; + let (writer, mut peer) = tokio::io::duplex(4096); + let mut transport = AsyncRwTransport::::new(&input[..], writer); + let write = transport.write.clone(); + let guard = write.lock().await; + { + let mut receive = Box::pin(transport.receive()); + assert!(futures::poll!(&mut receive).is_pending()); + } + drop(guard); + + let next = transport + .receive() + .await + .expect("the next request survives"); + assert_eq!(serde_json::to_value(next).unwrap()["id"], 1); + transport.close().await.unwrap(); + + let mut bytes = Vec::new(); + peer.read_to_end(&mut bytes).await.unwrap(); + let reply: serde_json::Value = serde_json::from_slice(&bytes) + .expect("cancelling a writer-lock wait must not lose the protocol error reply"); + assert_eq!( + reply, + serde_json::json!({ + "jsonrpc": "2.0", + "error": {"code": -32600, "message": "Invalid request"}, + }), + ); + } + + #[cfg(feature = "server")] + #[tokio::test] + async fn protocol_error_reply_survives_cancellation_during_flush() { + use tokio::io::AsyncReadExt; + + use crate::RoleServer; + + let input = b"{\"foo\":\"bar\"}\n{\"jsonrpc\":\"2.0\",\"method\":\"ping\",\"id\":1}\n"; + // A single byte of capacity suspends flushing after only part of the reply. + let (writer, mut peer) = tokio::io::duplex(1); + let mut transport = AsyncRwTransport::::new(&input[..], writer); + { + let mut receive = Box::pin(transport.receive()); + assert!(futures::poll!(&mut receive).is_pending()); + } + + let output = tokio::spawn(async move { + let mut bytes = Vec::new(); + peer.read_to_end(&mut bytes).await.unwrap(); + bytes + }); + let next = transport + .receive() + .await + .expect("the next request survives"); + assert_eq!(serde_json::to_value(next).unwrap()["id"], 1); + transport.close().await.unwrap(); + + let bytes = output.await.unwrap(); + let reply: serde_json::Value = serde_json::from_slice(&bytes) + .expect("cancelling a partial flush must preserve one complete protocol error reply"); + assert_eq!( + reply, + serde_json::json!({ + "jsonrpc": "2.0", + "error": {"code": -32600, "message": "Invalid request"}, + }), + ); + } }