fix(mcp): keep the in-flight reply when answering a server request fails

The reply to a server request (`ping`, or -32601 for anything else) was
sent with `self.send(reply).await?`. A failed send failed the
`tools/call` it interrupted, even though the server's real answer was
still on its way, and marked the connection disconnected. Over HTTP it
was worse: a 404/405/406/415/501 on that POST made the transport fall
back to legacy SSE, dropping the Streamable HTTP messages still queued
for the request in flight.

The reply now goes through a new `McpTransport::send_reply`, which
defaults to `send`; HttpTransport overrides it to report a refusal
without changing mode. A failure is logged at debug level and not
propagated, and the connection state is left alone.

Tests: cargo test -p codewhale-engine --lib mcp:: (373 passed); clippy
--all-targets -D warnings clean. The two new tests fail without this
change.
This commit is contained in:
luochen211 committed 2026-11-12 13:47:27 +08:00
1 parent a43de884ca
commit 8e8cd373ad
3 files changed
+143 -1

No files matched your search

+29 -1
View File
@@ -1669,6 +1669,10 @@ pub trait McpTransport: Send + Sync {
async fn recv(&mut self) -> Result<Vec<u8>>;
async fn send_reply(&mut self, msg: Vec<u8>) -> Result<()> {
self.send(msg).await
}
/// Record the protocol revision negotiated at `initialize`. Only the
/// Streamable HTTP transport uses it (the `MCP-Protocol-Version` header on
/// subsequent requests); stdio and legacy SSE have no header channel, so
@@ -3125,6 +3129,30 @@ impl McpConnection {
result
}
async fn answer_server_request(&mut self, reply: &serde_json::Value) {
let result = match serde_json::to_vec(reply) {
Ok(bytes) => {
let cancel_token = self.cancel_token.clone();
tokio::select! {
biased;
_ = cancel_token.cancelled() => {
Err(anyhow::anyhow!("MCP connection was cancelled"))
}
result = self.transport.send_reply(bytes) => result,
}
}
Err(err) => Err(err.into()),
};
if let Err(err) = result {
tracing::debug!(
target: "mcp",
server = %self.name,
error = %err,
"could not answer an MCP server request"
);
}
}
/// Handshake and discovery receive: each frame wait is bounded by the read
/// knob, and a server that stays silent past it is treated as dead.
async fn recv(&mut self, expected_id: String) -> Result<serde_json::Value> {
@@ -3214,7 +3242,7 @@ impl McpConnection {
"error": {"code": -32601, "message": format!("Method not found: {method}")}
})
};
self.send(reply).await?;
self.answer_server_request(&reply).await;
}
continue;
}
+15
View File
@@ -177,6 +177,21 @@ impl McpTransport for HttpTransport {
}
}
async fn send_reply(&mut self, msg: Vec<u8>) -> Result<()> {
match &mut self.mode {
HttpTransportMode::Streamable(transport) => {
transport.send(msg).await.map_err(|err| match err {
StreamableSendError::Incompatible(detail)
| StreamableSendError::StaleSession(detail) => anyhow::anyhow!(
"MCP Streamable HTTP refused a reply to a server request: {detail}"
),
StreamableSendError::Other(err) => err,
})
}
HttpTransportMode::Sse(transport) => transport.send(msg).await,
}
}
async fn recv(&mut self) -> Result<Vec<u8>> {
match &mut self.mode {
HttpTransportMode::Streamable(transport) => transport.recv().await,
+99
View File
@@ -3343,6 +3343,51 @@ async fn call_method_answers_server_requests_instead_of_taking_them_as_the_reply
assert_eq!(sent[2]["error"]["code"], -32601);
}
struct ReplyRefusingTransport {
sent: Arc<Mutex<Vec<serde_json::Value>>>,
responses: VecDeque<Vec<u8>>,
}
#[async_trait::async_trait]
impl McpTransport for ReplyRefusingTransport {
async fn send(&mut self, msg: Vec<u8>) -> Result<()> {
let msg: serde_json::Value = serde_json::from_slice(&msg)?;
anyhow::ensure!(msg.get("method").is_some(), "status=405 Method Not Allowed");
self.sent.lock().unwrap().push(msg);
Ok(())
}
async fn recv(&mut self) -> Result<Vec<u8>> {
self.responses
.pop_front()
.context("scripted transport exhausted")
}
}
#[tokio::test]
async fn call_method_keeps_its_reply_when_answering_a_server_request_fails() {
let sent = Arc::new(Mutex::new(Vec::new()));
let transport = ReplyRefusingTransport {
sent: Arc::clone(&sent),
responses: VecDeque::from([
json_frame(serde_json::json!({"jsonrpc": "2.0", "id": 1, "method": "ping"})),
json_frame(serde_json::json!({"jsonrpc": "2.0", "id": 1, "result": {"ok": true}})),
]),
};
let mut conn = test_connection(Box::new(transport));
let result = conn
.call_method("tools/call", serde_json::json!({"name": "echo"}), 1)
.await
.unwrap();
assert_eq!(result, serde_json::json!({"ok": true}));
assert_eq!(conn.state, ConnectionState::Ready);
let sent = sent.lock().unwrap();
assert_eq!(sent.len(), 1);
assert_eq!(sent[0]["method"], "tools/call");
}
#[tokio::test]
async fn call_method_invalid_json_includes_server_output_preview() {
let sent = Arc::new(Mutex::new(Vec::new()));
@@ -6799,6 +6844,60 @@ async fn custom_headers_applied_to_get_preflight() {
);
}
#[tokio::test]
async fn refused_reply_to_a_server_request_does_not_fall_back_to_sse() {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
let _env = crate::test_support::lock_test_env();
let _lock = lock_mcp_loopback_tests().await;
let _no_proxy = crate::test_support::EnvVarGuard::set("NO_PROXY", "*,127.0.0.1");
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let requests = Arc::new(Mutex::new(Vec::new()));
let requests_srv = Arc::clone(&requests);
let server = tokio::spawn(async move {
loop {
let (mut socket, _) = listener.accept().await.unwrap();
let mut buf = [0u8; 4096];
let n = socket.read(&mut buf).await.unwrap();
let request_line = String::from_utf8_lossy(&buf[..n])
.lines()
.next()
.unwrap_or_default()
.to_string();
requests_srv.lock().unwrap().push(request_line);
socket
.write_all(
b"HTTP/1.1 405 Method Not Allowed\r\nConnection: close\r\nContent-Length: 0\r\n\r\n",
)
.await
.unwrap();
}
});
let url = format!("http://{addr}/mcp");
let mut transport = HttpTransport::new(
test_mcp_http_client(&url),
url,
tokio_util::sync::CancellationToken::new(),
Duration::from_secs(2),
);
let result = transport
.send_reply(json_frame(
serde_json::json!({"jsonrpc": "2.0", "id": 1, "result": {}}),
))
.await;
server.abort();
assert!(result.is_err(), "a refused reply reports the refusal");
let requests = requests.lock().unwrap();
assert_eq!(requests.len(), 1, "requests seen: {requests:?}");
assert!(requests[0].starts_with("POST "), "{requests:?}");
}
// === add_runtime_server_config conflict tests ===
#[test]