Skip to content

Commit 398377a

Browse files
committed
feat: support ping requests before initialize handshake
1 parent 3d2c951 commit 398377a

File tree

2 files changed

+70
-18
lines changed

2 files changed

+70
-18
lines changed

crates/rmcp/src/service/server.rs

Lines changed: 29 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -131,22 +131,6 @@ where
131131
.ok_or_else(|| ServerInitializeError::ConnectionClosed(context.to_string()))
132132
}
133133

134-
/// Helper function to expect a request from the stream
135-
async fn expect_request<T>(
136-
transport: &mut T,
137-
context: &str,
138-
) -> Result<(ClientRequest, RequestId), ServerInitializeError>
139-
where
140-
T: Transport<RoleServer>,
141-
{
142-
let msg = expect_next_message(transport, context).await?;
143-
let msg_clone = msg.clone();
144-
msg.into_request()
145-
.ok_or(ServerInitializeError::ExpectedInitializeRequest(Some(
146-
msg_clone,
147-
)))
148-
}
149-
150134
pub async fn serve_server_with_ct<S, T, E, A>(
151135
service: S,
152136
transport: T,
@@ -177,8 +161,35 @@ where
177161
let mut transport = transport.into_transport();
178162
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
179163

180-
// Get initialize request
181-
let (request, id) = expect_request(&mut transport, "initialized request").await?;
164+
// Get initialize request; the MCP spec permits ping before initialize.
165+
// See: https://modelcontextprotocol.io/specification/2025-11-25/basic/lifecycle#initialization
166+
let (request, id) = loop {
167+
let msg = expect_next_message(&mut transport, "initialize request").await?;
168+
match msg {
169+
ClientJsonRpcMessage::Request(req)
170+
if matches!(req.request, ClientRequest::PingRequest(_)) =>
171+
{
172+
transport
173+
.send(ServerJsonRpcMessage::response(
174+
ServerResult::EmptyResult(EmptyResult {}),
175+
req.id,
176+
))
177+
.await
178+
.map_err(|error| {
179+
ServerInitializeError::transport::<T>(
180+
error,
181+
"sending pre-init ping response",
182+
)
183+
})?;
184+
}
185+
ClientJsonRpcMessage::Request(req) => break (req.request, req.id),
186+
other => {
187+
return Err(ServerInitializeError::ExpectedInitializeRequest(Some(
188+
other,
189+
)));
190+
}
191+
}
192+
};
182193

183194
let ClientRequest::InitializeRequest(peer_info) = &request else {
184195
return Err(ServerInitializeError::ExpectedInitializeRequest(Some(

crates/rmcp/tests/test_server_initialization.rs

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,47 @@ async fn server_init_succeeds_after_set_level_before_initialized() {
9696
result.unwrap().cancel().await.unwrap();
9797
}
9898

99+
// Server responds with EmptyResult to ping received before initialize request.
100+
#[tokio::test]
101+
async fn server_init_ping_response_is_empty_result_before_initialize() {
102+
let (server_transport, client_transport) = tokio::io::duplex(4096);
103+
let _server = tokio::spawn(async move { TestServer::new().serve(server_transport).await });
104+
let mut client = IntoTransport::<rmcp::RoleClient, _, _>::into_transport(client_transport);
105+
106+
client.send(ping_request(1)).await.unwrap();
107+
108+
let response = client.receive().await.unwrap();
109+
assert!(
110+
matches!(
111+
response,
112+
ServerJsonRpcMessage::Response(ref r)
113+
if matches!(r.result, ServerResult::EmptyResult(_))
114+
),
115+
"expected EmptyResult for pre-initialize ping, got: {response:?}"
116+
);
117+
}
118+
119+
// Server initializes successfully when ping is sent before the initialize request.
120+
#[tokio::test]
121+
async fn server_init_succeeds_after_ping_before_initialize() {
122+
let (server_transport, client_transport) = tokio::io::duplex(4096);
123+
let server_handle =
124+
tokio::spawn(async move { TestServer::new().serve(server_transport).await });
125+
let mut client = IntoTransport::<rmcp::RoleClient, _, _>::into_transport(client_transport);
126+
127+
client.send(ping_request(1)).await.unwrap();
128+
let _pong = client.receive().await.unwrap();
129+
do_initialize(&mut client).await;
130+
client.send(initialized_notification()).await.unwrap();
131+
132+
let result = server_handle.await.unwrap();
133+
assert!(
134+
result.is_ok(),
135+
"server should initialize successfully after pre-initialize ping"
136+
);
137+
result.unwrap().cancel().await.unwrap();
138+
}
139+
99140
// Server responds with EmptyResult to ping received before initialized.
100141
#[tokio::test]
101142
async fn server_init_ping_response_is_empty_result() {

0 commit comments

Comments
 (0)