Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions crates/rmcp/src/service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1103,9 +1103,15 @@ where
JsonRpcMessage::Error(error) => error.id.as_ref(),
_ => None,
} {
if let Some(ct) = local_ct_pool.remove(id) {
ct.cancel();
}
// A response whose request is no longer tracked means the
// request was cancelled (the cancelled-notification handler
// removed it from the pool). Per the spec, a receiver SHOULD
// NOT send a response for a cancelled request, so drop it.
let Some(ct) = local_ct_pool.remove(id) else {
tracing::debug!(%id, "dropping response for cancelled request");
continue;
};
ct.cancel();
let send = transport.send(m);
let current_span = tracing::Span::current();
response_send_tasks.spawn(async move {
Expand Down
211 changes: 211 additions & 0 deletions crates/rmcp/tests/test_cancelled_response.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,211 @@
//! A receiver SHOULD NOT send a response for a request it has already been told
//! to cancel. This drives a real stdio server with raw JSON-RPC: the tool blocks
//! until the request is cancelled, so its result is only produced *after* the
//! cancellation — the service loop must drop it rather than write it to the wire.

use std::{collections::BTreeSet, process::Stdio, time::Duration};

use rmcp::{
ErrorData as McpError, RoleServer, ServerHandler, ServiceExt,
model::{CallToolRequestParams, CallToolResult, ContentBlock, ServerCapabilities, ServerInfo},
service::RequestContext,
};
use serde_json::{Value, json};
use tokio::{
io::{AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader},
process::{Child, Command},
};

const HELPER_ENV: &str = "RMCP_CANCELLED_RESPONSE_HELPER";
const READ_TIMEOUT: Duration = Duration::from_secs(10);

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn cancelled_request_receives_no_response() -> anyhow::Result<()> {
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);

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?;
collect_ids_until(&mut reader, 1, READ_TIMEOUT).await?;
send_json(
&mut writer,
&json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }),
)
.await?;

// Start a request that blocks until cancelled, then cancel it. Its response is
// produced only after the cancellation arrives, so it must be suppressed.
send_json(
&mut writer,
&json!({
"jsonrpc": "2.0",
"id": 2,
"method": "tools/call",
"params": { "name": "wait-for-cancel", "arguments": {} }
}),
)
.await?;
send_json(
&mut writer,
&json!({
"jsonrpc": "2.0",
"method": "notifications/cancelled",
"params": { "requestId": 2 }
}),
)
.await?;
// A ping proves the server is alive past the cancellation, so the absence of
// an id=2 response is genuine suppression rather than a dead connection.
send_json(
&mut writer,
&json!({ "jsonrpc": "2.0", "id": 3, "method": "ping" }),
)
.await?;

let seen = collect_ids_until(&mut reader, 3, READ_TIMEOUT).await?;
assert!(seen.contains(&3));
assert!(!seen.contains(&2));

drop(writer);
wait_for_child(&mut child).await;
Ok(())
}

struct WaitForCancelServer;

impl ServerHandler for WaitForCancelServer {
fn get_info(&self) -> ServerInfo {
ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
}

async fn call_tool(
&self,
_request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> Result<CallToolResult, McpError> {
context.ct.cancelled().await;
Ok(CallToolResult::success(vec![ContentBlock::text(
"late response",
)]))
}
}

#[tokio::test]
async fn cancelled_response_helper() -> anyhow::Result<()> {
if std::env::var(HELPER_ENV).as_deref() != Ok("1") {
return Ok(());
}
run_helper_server().await?;
Ok(())
}

#[cfg(feature = "local")]
async fn run_helper_server() -> anyhow::Result<()> {
tokio::task::LocalSet::new()
.run_until(serve_helper_stdio())
.await
}

#[cfg(not(feature = "local"))]
async fn run_helper_server() -> anyhow::Result<()> {
serve_helper_stdio().await
}

async fn serve_helper_stdio() -> anyhow::Result<()> {
let server = WaitForCancelServer.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("cancelled_response_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;
}
}

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(())
}

/// Read response lines, collecting every message id seen, until `stop_id` is seen
/// (then a short grace read to catch any straggler) or the timeout elapses.
async fn collect_ids_until<R>(
reader: &mut BufReader<R>,
stop_id: u64,
timeout: Duration,
) -> anyhow::Result<BTreeSet<u64>>
where
R: tokio::io::AsyncRead + Unpin,
{
let mut seen = BTreeSet::new();
let mut deadline = tokio::time::Instant::now() + timeout;
loop {
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;
};
if read_result? == 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) {
seen.insert(id);
if id == stop_id {
// Give any late (incorrectly-sent) response a brief window to arrive.
deadline = tokio::time::Instant::now() + Duration::from_millis(300);
}
}
}
Ok(seen)
}