fix(transport): make AsyncRwTransport::receive cancel-safe (#941)

crates/rmcp/src/transport/async_rw.rs

@@ -124,22 +124,41 @@ where
 
     async fn receive(&mut self) -> Option<RxJsonRpcMessage<Role>> {
         loop {
-            self.line_buf.clear();
+            // `read_until` is not cancellation-safe on its own, and `receive` is
+            // polled inside a `select!` in the service loop: an in-progress line
+            // read is dropped whenever another branch (e.g. an outgoing response)
+            // becomes ready. We rely on `read_until` appending into `self.line_buf`
+            // and only returning at a delimiter or EOF, so a cancelled read leaves
+            // its partial bytes in `self.line_buf`. Keeping that buffer across
+            // calls lets the next read resume the same line; it is cleared only
+            // after a whole line has been consumed. Clearing at the top of the
+            // loop (the previous behaviour) discarded the partial read and so
+            // dropped incoming requests under concurrent response load.
             match self.read.read_until(b'\n', &mut self.line_buf).await {
+                // EOF. Any bytes still in `line_buf` are an incomplete trailing
+                // message with no delimiter, so there is nothing to deliver.
                 Ok(0) => return None,
                 Ok(_) => {}
                 Err(e) => {
                     tracing::error!("Error reading from stream: {}", e);
                     return None;
                 }
             }
-            let line = without_carriage_return(
-                self.line_buf.strip_suffix(b"\n").unwrap_or(&self.line_buf),
-            );
-            if line.is_empty() {
-                continue;
-            }
-            match try_parse_with_compatibility::<RxJsonRpcMessage<Role>>(line, "receive") {
+            // A returned `read_until` means a full line is buffered. Parse it
+            // (borrowing `line_buf`), then clear the buffer — retaining its
+            // capacity for the next read — before handling the parse result.
+            let parsed = {
+                let line = without_carriage_return(
+                    self.line_buf.strip_suffix(b"\n").unwrap_or(&self.line_buf),
+                );
+                if line.is_empty() {
+                    self.line_buf.clear();
+                    continue;
+                }
+                try_parse_with_compatibility::<RxJsonRpcMessage<Role>>(line, "receive")
+            };
+            self.line_buf.clear();
+            match parsed {
                 Ok(Some(msg)) => return Some(msg),
                 Ok(None) => continue,
                 Err(JsonRpcMessageCodecError::Serde(e)) => {

crates/rmcp/tests/test_stdio_response_concurrency.rs

@@ -0,0 +1,198 @@
+#![cfg(not(feature = "local"))]
+
+use std::{collections::BTreeSet, process::Stdio, time::Duration};
+
+use rmcp::{
+    ErrorData as McpError, ServerHandler, ServiceExt,
+    model::{CallToolRequestParams, CallToolResult, ContentBlock, ServerCapabilities, ServerInfo},
+};
+use serde_json::{Value, json};
+use tokio::{
+    io::{AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader},
+    process::{Child, Command},
+};
+
+const HELPER_ENV: &str = "RMCP_STDIO_RESPONSE_CONCURRENCY_HELPER";
+const REQUESTS: usize = 200;
+const RESPONSE_BYTES: usize = 64 * 1024;
+const READ_TIMEOUT: Duration = Duration::from_secs(10);
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
+async fn raw_client_concurrent_large_stdio_tool_responses_are_not_lost() -> anyhow::Result<()> {
+    // Spawn the same test binary as a child process so the server uses real
+    // stdio pipes, not an in-process transport.
+    let mut child = spawn_helper();
+    let mut writer = child.stdin.take().expect("helper stdin");
+    let stdout = child.stdout.take().expect("helper stdout");
+    let mut reader = BufReader::new(stdout);
+
+    // Complete the normal MCP initialization flow before stressing tools/call.
+    send_json(
+        &mut writer,
+        &json!({
+            "jsonrpc": "2.0",
+            "id": 1,
+            "method": "initialize",
+            "params": {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {},
+                "clientInfo": { "name": "raw-test-client", "version": "0.0.0" }
+            }
+        }),
+    )
+    .await?;
+    read_response_for_id(&mut reader, 1).await?;
+
+    send_json(
+        &mut writer,
+        &json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }),
+    )
+    .await?;
+
+    // Send the whole batch before reading responses. This creates concurrent
+    // request handling and concurrent response production inside rmcp.
+    for id in request_ids() {
+        send_json(
+            &mut writer,
+            &json!({
+                "jsonrpc": "2.0",
+                "id": id,
+                "method": "tools/call",
+                "params": { "name": "large-response", "arguments": {} }
+            }),
+        )
+        .await?;
+    }
+
+    let missing_ids = read_responses_for_ids(&mut reader, request_ids(), READ_TIMEOUT).await?;
+    assert!(
+        missing_ids.is_empty(),
+        "missing response ids: {missing_ids:?}"
+    );
+
+    drop(writer);
+    wait_for_child(&mut child).await;
+    Ok(())
+}
+
+struct LargeResponseServer;
+
+impl ServerHandler for LargeResponseServer {
+    #[allow(deprecated)]
+    fn get_info(&self) -> ServerInfo {
+        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
+    }
+
+    async fn call_tool(
+        &self,
+        request: CallToolRequestParams,
+        _context: rmcp::service::RequestContext<rmcp::RoleServer>,
+    ) -> Result<CallToolResult, McpError> {
+        assert_eq!("large-response", request.name.as_ref());
+        Ok(CallToolResult::success(vec![ContentBlock::text(
+            "x".repeat(RESPONSE_BYTES),
+        )]))
+    }
+}
+
+#[tokio::test]
+async fn stdio_response_concurrency_helper() -> anyhow::Result<()> {
+    // The parent test starts this same binary with HELPER_ENV=1 so it can act
+    // as a small MCP server connected over real stdin/stdout pipes.
+    if std::env::var(HELPER_ENV).as_deref() != Ok("1") {
+        return Ok(());
+    }
+    let server = LargeResponseServer.serve(rmcp::transport::stdio()).await?;
+    server.waiting().await?;
+    Ok(())
+}
+
+fn spawn_helper() -> Child {
+    let exe = std::env::current_exe().expect("current test exe");
+    Command::new(exe)
+        .arg("--exact")
+        .arg("stdio_response_concurrency_helper")
+        .arg("--quiet")
+        .arg("--nocapture")
+        .arg("--test-threads")
+        .arg("1")
+        .env(HELPER_ENV, "1")
+        .stdin(Stdio::piped())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::null())
+        .kill_on_drop(true)
+        .spawn()
+        .expect("spawn helper")
+}
+
+async fn wait_for_child(child: &mut Child) {
+    let _ = tokio::time::timeout(Duration::from_secs(2), child.wait()).await;
+    if child.id().is_some() {
+        let _ = child.kill().await;
+    }
+}
+
+fn request_ids() -> BTreeSet<u64> {
+    (1000..1000 + REQUESTS as u64).collect()
+}
+
+async fn send_json<W>(writer: &mut W, message: &Value) -> anyhow::Result<()>
+where
+    W: AsyncWrite + Unpin,
+{
+    let serialized = serde_json::to_string(message)?;
+    writer.write_all(serialized.as_bytes()).await?;
+    writer.write_all(b"\n").await?;
+    writer.flush().await?;
+    Ok(())
+}
+
+async fn read_response_for_id<R>(reader: &mut BufReader<R>, expected_id: u64) -> anyhow::Result<()>
+where
+    R: tokio::io::AsyncRead + Unpin,
+{
+    let missing =
+        read_responses_for_ids(reader, BTreeSet::from([expected_id]), READ_TIMEOUT).await?;
+    if missing.is_empty() {
+        Ok(())
+    } else {
+        anyhow::bail!("missing response id {expected_id}")
+    }
+}
+
+async fn read_responses_for_ids<R>(
+    reader: &mut BufReader<R>,
+    mut pending_ids: BTreeSet<u64>,
+    timeout: Duration,
+) -> anyhow::Result<BTreeSet<u64>>
+where
+    R: tokio::io::AsyncRead + Unpin,
+{
+    let deadline = tokio::time::Instant::now() + timeout;
+    while !pending_ids.is_empty() {
+        let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
+        if remaining.is_zero() {
+            break;
+        }
+        let mut line = String::new();
+        let Ok(read_result) = tokio::time::timeout(remaining, reader.read_line(&mut line)).await
+        else {
+            break;
+        };
+        let read = read_result?;
+        if read == 0 {
+            break;
+        }
+        let trimmed = line.trim();
+        if trimmed.is_empty() {
+            continue;
+        }
+        let Ok(value) = serde_json::from_str::<Value>(trimmed) else {
+            continue;
+        };
+        if let Some(id) = value.get("id").and_then(Value::as_u64) {
+            pending_ids.remove(&id);
+        }
+    }
+    Ok(pending_ids)
+}