Skip to content

Commit a156be8

Browse files
gocamillegocamille
authored
fix(client): answer a server ping during the initialize handshake (#1366)
* fix(client): answer a server ping during the initialize handshake expect_response dropped a server ping while waiting for the initialize response, so the ping never got a reply. Answer it with an empty result during the initialize wait, the same way the server answers an early client ping. The server/discover wait still leaves it unanswered. Refs #745 Fixes #1362 * test(client): drop duplicated discover ping comment * refactor(client): replace answer_ping bool with PreInitPing enum * test(client): assert notifications/initialized follows the ping response * refactor(client): match on PreInitPing instead of comparing it The compiler now flags this spot if PreInitPing gains a variant, and the PartialEq and Eq derives are no longer needed.
1 parent b7568b8 commit a156be8

3 files changed

Lines changed: 165 additions & 10 deletions

File tree

‎crates/rmcp/src/service/client.rs‎

Lines changed: 51 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -158,22 +158,35 @@ where
158158
.ok_or_else(|| ClientInitializeError::ConnectionClosed(context.to_string()))
159159
}
160160

161+
/// Whether `expect_response` answers a server ping that arrives before the
162+
/// handshake completes.
163+
#[derive(Clone, Copy)]
164+
enum PreInitPing {
165+
Answer,
166+
Ignore,
167+
}
168+
161169
/// Helper function to expect a response from the stream, correlated to
162170
/// `expected_id`.
163171
///
164172
/// Both success and error responses are checked here: a mismatched id on a
165173
/// success response is `ConflictInitResponseId`; on an error response (whose
166174
/// `id` is optional per spec) it is `UncorrelatedErrorResponse`. The caller
167175
/// never sees an uncorrelated response.
176+
///
177+
/// A server ping is answered only with `PreInitPing::Answer`. The discover
178+
/// wait, including Auto's timeout-bounded probe, leaves it unanswered:
179+
/// 2026-07-28 has no ping, and a cancelled send could leave a partial frame.
168180
async fn expect_response<T, S>(
169181
transport: &mut T,
170182
context: &str,
171183
service: &S,
172184
peer: Peer<RoleClient>,
173185
expected_id: &RequestId,
186+
pre_init_ping: PreInitPing,
174187
) -> Result<ServerResult, ClientInitializeError>
175188
where
176-
T: Transport<RoleClient>,
189+
T: Transport<RoleClient> + 'static,
177190
S: Service<RoleClient>,
178191
{
179192
loop {
@@ -231,10 +244,26 @@ where
231244
}
232245
}
233246
// Server could send pings before handshake
234-
ServerJsonRpcMessage::Request(ref request)
247+
ServerJsonRpcMessage::Request(request)
235248
if matches!(request.request, ServerRequest::PingRequest(_)) =>
236249
{
237-
tracing::trace!("Received ping request. Ignored.")
250+
match pre_init_ping {
251+
PreInitPing::Answer => {
252+
transport
253+
.send(ClientJsonRpcMessage::response(
254+
ClientResult::empty(()),
255+
request.id,
256+
))
257+
.await
258+
.map_err(|error| {
259+
ClientInitializeError::transport::<T>(
260+
error,
261+
"send pre-init ping response",
262+
)
263+
})?;
264+
}
265+
PreInitPing::Ignore => tracing::trace!("Received ping request. Ignored."),
266+
}
238267
}
239268
// Server SHOULD NOT send any other messages before handshake. We ignore them anyway
240269
_ => tracing::warn!(?message, "Received unexpected message"),
@@ -948,8 +977,15 @@ where
948977
context: "send initialize request".into(),
949978
})?;
950979

951-
let response =
952-
expect_response(transport, "initialize response", service, peer.clone(), &id).await?;
980+
let response = expect_response(
981+
transport,
982+
"initialize response",
983+
service,
984+
peer.clone(),
985+
&id,
986+
PreInitPing::Answer,
987+
)
988+
.await?;
953989

954990
let ServerResult::InitializeResult(initialize_result) = response else {
955991
return Err(ClientInitializeError::ExpectedInitResult(Some(response)));
@@ -1019,7 +1055,16 @@ where
10191055
ClientInitializeError::transport::<T>(error, "send discover request")
10201056
})?;
10211057

1022-
match expect_response(transport, "discover response", service, peer.clone(), &id).await {
1058+
match expect_response(
1059+
transport,
1060+
"discover response",
1061+
service,
1062+
peer.clone(),
1063+
&id,
1064+
PreInitPing::Ignore,
1065+
)
1066+
.await
1067+
{
10231068
Ok(ServerResult::DiscoverResult(result)) => {
10241069
let Some(selected) =
10251070
select_protocol_version(&preferred_versions, &result.supported_versions)

‎crates/rmcp/tests/test_client_initialization.rs‎

Lines changed: 53 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,9 @@ use common::handlers::TestClientHandler;
99
use rmcp::{
1010
ServiceExt,
1111
model::{
12-
ClientJsonRpcMessage, ErrorCode, ErrorData, InitializeResult, JsonRpcError,
13-
JsonRpcVersion2_0, RequestId, ServerCapabilities, ServerJsonRpcMessage, ServerResult,
12+
ClientConfig, ClientJsonRpcMessage, ClientNotification, ClientResult, ErrorCode, ErrorData,
13+
InitializeResult, JsonRpcError, JsonRpcVersion2_0, PingRequest, ProtocolVersion, RequestId,
14+
ServerCapabilities, ServerJsonRpcMessage, ServerRequest, ServerResult,
1415
},
1516
transport::{IntoTransport, Transport},
1617
};
@@ -157,3 +158,53 @@ async fn test_client_init_handles_jsonrpc_error() {
157158
_ => panic!("Expected ClientInitializeError::JsonRpcError"),
158159
}
159160
}
161+
162+
#[tokio::test]
163+
async fn client_answers_server_ping_during_initialize() {
164+
let (server_transport, client_transport) = tokio::io::duplex(4096);
165+
let mut server = IntoTransport::<rmcp::RoleServer, _, _>::into_transport(server_transport);
166+
let server_task = tokio::spawn(async move {
167+
let ClientJsonRpcMessage::Request(initialize) =
168+
server.receive().await.expect("expected initialize request")
169+
else {
170+
panic!("expected initialize request");
171+
};
172+
let ping_id = RequestId::String("server-ping".into());
173+
server
174+
.send(ServerJsonRpcMessage::request(
175+
ServerRequest::PingRequest(PingRequest::default()),
176+
ping_id.clone(),
177+
))
178+
.await
179+
.expect("send ping");
180+
server
181+
.send(ServerJsonRpcMessage::response(
182+
ServerResult::InitializeResult(
183+
InitializeResult::new(ServerCapabilities::default())
184+
.with_protocol_version(ProtocolVersion::LATEST_WITH_INITIALIZE),
185+
),
186+
initialize.id,
187+
))
188+
.await
189+
.expect("send initialize response");
190+
191+
let Some(ClientJsonRpcMessage::Response(response)) = server.receive().await else {
192+
panic!("expected the ping response before notifications/initialized");
193+
};
194+
assert_eq!(response.id, ping_id);
195+
assert!(matches!(response.result, ClientResult::EmptyResult(_)));
196+
assert!(matches!(
197+
server.receive().await,
198+
Some(ClientJsonRpcMessage::Notification(notification))
199+
if matches!(notification.notification, ClientNotification::InitializedNotification(_))
200+
));
201+
});
202+
203+
let client = ClientConfig::default()
204+
.with_protocol_version(ProtocolVersion::LATEST_WITH_INITIALIZE)
205+
.serve(client_transport)
206+
.await
207+
.expect("initialize client");
208+
client.cancel().await.expect("cancel client");
209+
server_task.await.expect("server task");
210+
}

‎crates/rmcp/tests/test_client_lifecycle_modes.rs‎

Lines changed: 61 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@ use rmcp::{
44
ClientHandler, ClientLifecycleMode, ClientServiceExt, ServerHandler, ServiceExt,
55
model::{
66
ClientJsonRpcMessage, ClientRequest, DiscoverResult, ErrorCode, ErrorData, GetMeta,
7-
Implementation, InitializeResult, ProtocolVersion, RequestId, ServerCapabilities,
8-
ServerJsonRpcMessage, ServerResult,
7+
Implementation, InitializeResult, PingRequest, ProtocolVersion, RequestId,
8+
ServerCapabilities, ServerJsonRpcMessage, ServerRequest, ServerResult,
99
},
1010
service::PeerRequestOptions,
1111
transport::{IntoTransport, Transport},
@@ -277,6 +277,65 @@ async fn discover_startup_omits_initialize() {
277277
server_task.await.expect("server task");
278278
}
279279

280+
#[tokio::test]
281+
async fn discover_startup_ignores_server_ping() {
282+
let (server_transport, client_transport) = tokio::io::duplex(4096);
283+
let mut server = IntoTransport::<rmcp::RoleServer, _, _>::into_transport(server_transport);
284+
let server_task = tokio::spawn(async move {
285+
let ClientJsonRpcMessage::Request(request) =
286+
server.receive().await.expect("expected discover request")
287+
else {
288+
panic!("expected discover request");
289+
};
290+
server
291+
.send(ServerJsonRpcMessage::request(
292+
ServerRequest::PingRequest(PingRequest::default()),
293+
RequestId::String("server-ping".into()),
294+
))
295+
.await
296+
.expect("send ping");
297+
server
298+
.send(ServerJsonRpcMessage::response(
299+
ServerResult::DiscoverResult(DiscoverResult::new(
300+
vec![ProtocolVersion::V_2026_07_28],
301+
ServerCapabilities::default(),
302+
)),
303+
request.id,
304+
))
305+
.await
306+
.expect("send discover response");
307+
308+
// Anything written for the ping would arrive before this request.
309+
let Some(ClientJsonRpcMessage::Request(request)) = server.receive().await else {
310+
panic!("expected the next client request, not a ping response");
311+
};
312+
assert!(matches!(
313+
request.request,
314+
ClientRequest::ListToolsRequest(_)
315+
));
316+
server
317+
.send(ServerJsonRpcMessage::response(
318+
ServerResult::ListToolsResult(Default::default()),
319+
request.id,
320+
))
321+
.await
322+
.expect("send tools response");
323+
});
324+
325+
let client = DiscoverClient
326+
.serve_with_lifecycle(
327+
client_transport,
328+
ClientLifecycleMode::Discover {
329+
preferred_versions: vec![ProtocolVersion::V_2026_07_28],
330+
},
331+
)
332+
.await
333+
.expect("discover client should start");
334+
client.list_tools(None).await.expect("list tools");
335+
client.cancel().await.expect("cancel client");
336+
server_task.await.expect("server task");
337+
}
338+
280339
#[tokio::test]
281340
async fn auto_startup_falls_back_after_discover_method_not_found() {
282341
let (server_transport, client_transport) = tokio::io::duplex(4096);

0 commit comments

Comments
 (0)