From 287556208ad6b721e6161789e418afbbf368674a Mon Sep 17 00:00:00 2001 From: aniongithub Date: Mon, 7 Sep 2026 22:17:22 -0700 Subject: [PATCH] Fix MCP handshake failure with hosts that probe via server/discover The GitHub Copilot CLI (and similar hosts) send an optimistic `server/discover` request before the standard `initialize`. rmcp's handshake treats the first non-ping request as the initialize request; when it isn't, serve() returns an error and the process exits, closing stdin before the host's fallback `initialize` arrives (broken pipe), so the server never starts. Wrap the stdio transport in a DiscoveryGuard that replies to any pre-initialize request with JSON-RPC -32601 (Method not found) and keeps waiting, letting the host's fallback initialize succeed. This mirrors how spec-compliant servers (e.g. mind-map) already behave. Add end-to-end tests that replay the probe-then-initialize sequence over an in-memory transport and confirm a normal handshake is unaffected. --- .../devcontainer-mcp/src/discovery_guard.rs | 219 ++++++++++++++++++ crates/devcontainer-mcp/src/main.rs | 13 +- 2 files changed, 230 insertions(+), 2 deletions(-) create mode 100644 crates/devcontainer-mcp/src/discovery_guard.rs diff --git a/crates/devcontainer-mcp/src/discovery_guard.rs b/crates/devcontainer-mcp/src/discovery_guard.rs new file mode 100644 index 0000000..0862b51 --- /dev/null +++ b/crates/devcontainer-mcp/src/discovery_guard.rs @@ -0,0 +1,219 @@ +//! Transport wrapper that keeps the server alive through pre-initialize probes. +//! +//! Some MCP hosts (notably the GitHub Copilot CLI) send an optimistic +//! `server/discover` request *before* the standard `initialize` handshake. A +//! spec-compliant server is expected to answer an unknown method with a +//! JSON-RPC `-32601 Method not found` error and stay running, at which point +//! the host falls back to the normal `initialize` flow. +//! +//! rmcp's server handshake instead treats the first non-`ping` request as the +//! `initialize` request; when it isn't, `serve()` returns an error, `main` +//! exits, and stdin is closed before the host's fallback `initialize` arrives +//! ("broken pipe"). This wrapper intercepts any request received before +//! `initialize`, replies with `Method not found`, and keeps waiting so the +//! real handshake can proceed. + +use rmcp::model::{ + ClientJsonRpcMessage, ClientRequest, ErrorCode, ErrorData, RequestId, ServerJsonRpcMessage, +}; +use rmcp::service::RoleServer; +use rmcp::transport::Transport; + +pub struct DiscoveryGuard { + inner: T, + initialized: bool, +} + +impl DiscoveryGuard { + pub fn new(inner: T) -> Self { + Self { + inner, + initialized: false, + } + } +} + +enum PreInit { + Forward, + Reply(RequestId, String), +} + +impl Transport for DiscoveryGuard +where + T: Transport + Send, +{ + type Error = T::Error; + + fn send( + &mut self, + item: ServerJsonRpcMessage, + ) -> impl std::future::Future> + Send + 'static { + self.inner.send(item) + } + + async fn receive(&mut self) -> Option { + loop { + let msg = self.inner.receive().await?; + if self.initialized { + return Some(msg); + } + + let action = match &msg { + ClientJsonRpcMessage::Request(req) => match &req.request { + ClientRequest::InitializeRequest(_) => { + self.initialized = true; + PreInit::Forward + } + // Let rmcp answer pre-init pings itself (allowed by the spec). + ClientRequest::PingRequest(_) => PreInit::Forward, + other => PreInit::Reply(req.id.clone(), other.method().to_string()), + }, + _ => PreInit::Forward, + }; + + match action { + PreInit::Forward => return Some(msg), + PreInit::Reply(id, method) => { + let response = ServerJsonRpcMessage::error( + ErrorData::new( + ErrorCode::METHOD_NOT_FOUND, + format!("method not available before initialization: {method}"), + None, + ), + id, + ); + // Best-effort: if the client hung up, the next receive() ends the loop. + let _ = self.inner.send(response).await; + continue; + } + } + } + } + + fn close(&mut self) -> impl std::future::Future> + Send { + self.inner.close() + } +} + +#[cfg(test)] +mod tests { + use super::DiscoveryGuard; + use crate::tools::DevContainerMcp; + use rmcp::service::RoleServer; + use rmcp::transport::IntoTransport; + use rmcp::ServiceExt; + use serde_json::{json, Value}; + use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream, ReadHalf, WriteHalf}; + + /// A minimal raw MCP client speaking newline-delimited JSON-RPC over an + /// in-memory pipe — enough to replay a host handshake byte-for-byte, + /// including the non-standard `server/discover` probe that rmcp's own + /// client would never send. + struct RawClient { + writer: WriteHalf, + reader: tokio::io::Lines>>, + } + + impl RawClient { + async fn send(&mut self, msg: Value) { + let mut line = msg.to_string(); + line.push('\n'); + self.writer.write_all(line.as_bytes()).await.unwrap(); + self.writer.flush().await.unwrap(); + } + + async fn recv(&mut self) -> Value { + let line = self + .reader + .next_line() + .await + .expect("read line") + .expect("server closed stream unexpectedly"); + serde_json::from_str(&line).expect("valid JSON-RPC line") + } + } + + /// Spawn the real `DevContainerMcp` server behind `DiscoveryGuard` over an + /// in-memory duplex and return a raw client wired to the other end. + fn spawn_server() -> (RawClient, tokio::task::JoinHandle<()>) { + let (client_io, server_io) = tokio::io::duplex(64 * 1024); + + let handle = tokio::spawn(async move { + let service = DevContainerMcp::new(); + let transport = + DiscoveryGuard::new(IntoTransport::::into_transport(server_io)); + let running = service.serve(transport).await.expect("server handshake"); + let _ = running.waiting().await; + }); + + let (read, write) = tokio::io::split(client_io); + let client = RawClient { + writer: write, + reader: BufReader::new(read).lines(), + }; + (client, handle) + } + + fn initialize_msg(id: i64) -> Value { + json!({ + "jsonrpc": "2.0", "id": id, "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": { "name": "test-client", "version": "0.0.0" } + } + }) + } + + /// Regression test for the reported bug: the GitHub Copilot CLI sends an + /// optimistic `server/discover` request before `initialize`. The server must + /// answer it with `-32601 Method not found`, stay alive, and then complete a + /// normal handshake so tools can be listed. + #[tokio::test] + async fn survives_pre_initialize_discover_probe() { + let (mut client, handle) = spawn_server(); + + // 1. Optimistic probe before initialize — must NOT crash the server. + client + .send(json!({"jsonrpc": "2.0", "id": 0, "method": "server/discover", "params": {}})) + .await; + let discover_reply = client.recv().await; + assert_eq!(discover_reply["id"], 0); + assert_eq!(discover_reply["error"]["code"], -32601); + + // 2. Fallback initialize — the server is still running and handshakes. + client.send(initialize_msg(1)).await; + let init_reply = client.recv().await; + assert_eq!(init_reply["id"], 1); + assert_eq!(init_reply["result"]["serverInfo"]["name"], "devcontainer-mcp"); + + // 3. Complete the lifecycle and confirm tools are usable. + client + .send(json!({"jsonrpc": "2.0", "method": "notifications/initialized"})) + .await; + client + .send(json!({"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}})) + .await; + let tools_reply = client.recv().await; + assert_eq!(tools_reply["id"], 2); + assert!( + !tools_reply["result"]["tools"].as_array().unwrap().is_empty(), + "expected a non-empty tool list" + ); + + handle.abort(); + } + + /// A normal handshake (no probe) must be unaffected by the wrapper. + #[tokio::test] + async fn normal_handshake_still_works() { + let (mut client, handle) = spawn_server(); + + client.send(initialize_msg(1)).await; + let init_reply = client.recv().await; + assert_eq!(init_reply["id"], 1); + assert_eq!(init_reply["result"]["serverInfo"]["name"], "devcontainer-mcp"); + + handle.abort(); + } +} diff --git a/crates/devcontainer-mcp/src/main.rs b/crates/devcontainer-mcp/src/main.rs index 62b9d69..9007a8b 100644 --- a/crates/devcontainer-mcp/src/main.rs +++ b/crates/devcontainer-mcp/src/main.rs @@ -1,9 +1,13 @@ +mod discovery_guard; mod tools; use clap::{Parser, Subcommand}; -use rmcp::transport::stdio; +use rmcp::service::RoleServer; +use rmcp::transport::{stdio, IntoTransport}; use rmcp::ServiceExt; +use discovery_guard::DiscoveryGuard; + #[derive(Parser)] #[command(name = "devcontainer-mcp")] #[command(about = "MCP server and CLI for managing DevContainers")] @@ -32,7 +36,12 @@ async fn main() -> anyhow::Result<()> { Commands::Serve => { tracing::info!("Starting devcontainer-mcp MCP server over stdio"); let service = tools::DevContainerMcp::new(); - let server = service.serve(stdio()).await?; + // Wrap stdio so an optimistic pre-initialize `server/discover` probe + // (sent by some hosts, e.g. the GitHub Copilot CLI) is answered with + // `Method not found` instead of aborting the handshake. + let transport = + DiscoveryGuard::new(IntoTransport::::into_transport(stdio())); + let server = service.serve(transport).await?; server.waiting().await?; } }