Skip to content

fix(client): ignore late discovery replies after legacy fallback #1436

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
charliermarsh wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

fix(client): ignore late discovery replies after legacy fallback #1436

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
45 changes: 32 additions & 13 deletions crates/rmcp/src/service/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,8 @@ enum PreInitPing {
/// Both success and error responses are checked here: a mismatched id on a
/// success response is `ConflictInitResponseId`; on an error response (whose
/// `id` is optional per spec) it is `UncorrelatedErrorResponse`. The caller
/// never sees an uncorrelated response.
/// never sees an uncorrelated response. During timeout fallback, replies to
/// the abandoned discovery request are skipped while waiting for `initialize`.
///
/// A server ping is answered only with `PreInitPing::Answer`. The discover
/// wait, including Auto's timeout-bounded probe, leaves it unanswered:
Expand All @@ -187,15 +188,21 @@ async fn expect_response<T, S>(
peer: Peer<RoleClient>,
expected_id: &RequestId,
pre_init_ping: PreInitPing,
timed_out_probe: Option<&RequestId>,
) -> Result<ServerResult, ClientInitializeError>
where
T: Transport<RoleClient> + 'static,
S: Service<RoleClient>,
{
let is_timed_out_probe =
|id: &RequestId| timed_out_probe.is_some_and(|probe| probe.matches_response_id(id));
loop {
let message = expect_next_message(transport, context).await?;
match message {
ServerJsonRpcMessage::Response(JsonRpcResponse { id, result, .. }) => {
if is_timed_out_probe(&id) {
continue;
}
if !expected_id.matches_response_id(&id) {
return Err(ClientInitializeError::ConflictInitResponseId(
expected_id.clone(),
Expand All @@ -205,6 +212,9 @@ where
return Ok(result);
}
ServerJsonRpcMessage::Error(error) => {
if error.id.as_ref().is_some_and(is_timed_out_probe) {
continue;
}
return Err(match &error.id {
Some(id) if expected_id.matches_response_id(id) => {
ClientInitializeError::JsonRpcError(error.error)
Expand Down Expand Up @@ -837,19 +847,19 @@ where
T: Transport<RoleClient> + 'static,
{
let mut transport = transport.into_transport();
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
let (peer, peer_rx) = Peer::new(id_provider.clone(), None);
let (peer, peer_rx) = Peer::new(Arc::<AtomicU32RequestIdProvider>::default(), None);
let mut last_discover_id = None;
let client_info = service.get_info();

match lifecycle {
ClientLifecycleMode::Initialize => {
legacy_startup(&service, &mut transport, &id_provider, &peer, client_info).await?;
legacy_startup(&service, &mut transport, &peer, client_info, None).await?;
}
ClientLifecycleMode::Discover { preferred_versions } => {
match discover_startup(
&service,
&mut transport,
&id_provider,
&mut last_discover_id,
&peer,
&client_info,
preferred_versions,
Expand All @@ -870,7 +880,7 @@ where
discover_startup(
&service,
&mut transport,
&id_provider,
&mut last_discover_id,
&peer,
&client_info,
preferred_versions,
Expand All @@ -885,8 +895,7 @@ where
legacy_info.protocol_version = version;
}
if let Err(fallback_error) =
legacy_startup(&service, &mut transport, &id_provider, &peer, legacy_info)
.await
legacy_startup(&service, &mut transport, &peer, legacy_info, None).await
{
return Err(ClientInitializeError::LegacyFallbackFailed {
discover: discover_error,
Expand All @@ -900,8 +909,14 @@ where
if let Some(version) = legacy_version {
legacy_info.protocol_version = version;
}
legacy_startup(&service, &mut transport, &id_provider, &peer, legacy_info)
.await?;
legacy_startup(
&service,
&mut transport,
&peer,
legacy_info,
last_discover_id.as_ref(),
)
.await?;
}
}
}
Expand Down Expand Up @@ -954,15 +969,15 @@ enum DiscoverOutcome {
async fn legacy_startup<S, T>(
service: &S,
transport: &mut T,
id_provider: &Arc<AtomicU32RequestIdProvider>,
peer: &Peer<RoleClient>,
client_info: InitializeRequestParams,
timed_out_probe: Option<&RequestId>,
) -> Result<(), ClientInitializeError>
where
S: Service<RoleClient>,
T: Transport<RoleClient> + 'static,
{
let id = id_provider.next_request_id();
let id = peer.request_id_provider.next_request_id();
let client_info_for_meta = client_info.clone();
let init_request = InitializeRequest {
method: Default::default(),
Expand All @@ -987,6 +1002,7 @@ where
peer.clone(),
&id,
PreInitPing::Answer,
timed_out_probe,
)
.await?;

Expand Down Expand Up @@ -1022,7 +1038,7 @@ where
async fn discover_startup<S, T>(
service: &S,
transport: &mut T,
id_provider: &Arc<AtomicU32RequestIdProvider>,
last_discover_id: &mut Option<RequestId>,
peer: &Peer<RoleClient>,
client_info: &InitializeRequestParams,
preferred_versions: Vec<ProtocolVersion>,
Expand All @@ -1035,6 +1051,7 @@ where
return Err(ClientInitializeError::NoPreferredProtocolVersion);
}

let id_provider = &peer.request_id_provider;
let mut attempted = Vec::new();
let mut candidate = preferred_versions[0].clone();
loop {
Expand All @@ -1048,6 +1065,7 @@ where
let mut discover = DiscoverRequest::new(DiscoverRequestParams {});
discover.extensions.insert(meta);
let id = id_provider.next_request_id();
*last_discover_id = Some(id.clone());
transport
.send(ClientJsonRpcMessage::request(
ClientRequest::DiscoverRequest(discover),
Expand All @@ -1065,6 +1083,7 @@ where
peer.clone(),
&id,
PreInitPing::Ignore,
None,
)
.await
{
Expand Down
111 changes: 86 additions & 25 deletions crates/rmcp/tests/test_client_lifecycle_modes.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -336,11 +336,31 @@ async fn discover_startup_ignores_server_ping() {
server_task.await.expect("server task");
}

#[tokio::test]
async fn auto_startup_falls_back_after_discover_method_not_found() {
#[rstest::rstest]
#[case::prompt_rejection(false, false)]
#[case::late_rejection(true, false)]
#[case::late_success(true, true)]
#[tokio::test(start_paused = true)]
async fn auto_fallback_uses_initialize_reply(
#[case] late: bool,
#[case] success: bool,
#[values(false, true)] retry: bool,
) {
let (server_transport, client_transport) = tokio::io::duplex(4096);
let mut server = IntoTransport::<rmcp::RoleServer, _, _>::into_transport(server_transport);
let server_task = tokio::spawn(async move {
if retry {
let (request, id) = server.receive().await.unwrap().into_request().unwrap();
assert!(matches!(request, ClientRequest::DiscoverRequest(_)));
let error = ErrorData::unsupported_protocol_version(
ProtocolVersion::V_2026_07_28,
&[ProtocolVersion::V_2026_07_28],
);
server
.send(ServerJsonRpcMessage::error(error, Some(id)))
.await
.unwrap();
}
let ClientJsonRpcMessage::Request(discover) =
server.receive().await.expect("expected discover request")
else {
Expand All @@ -350,13 +370,24 @@ async fn auto_startup_falls_back_after_discover_method_not_found() {
discover.request,
ClientRequest::DiscoverRequest(_)
));
server
.send(ServerJsonRpcMessage::error(
if late {
tokio::time::sleep(std::time::Duration::from_secs(12)).await;
}
let reply = if success {
ServerJsonRpcMessage::response(
ServerResult::DiscoverResult(DiscoverResult::new(
vec![ProtocolVersion::V_2026_07_28],
Default::default(),
)),
RequestId::String(discover.id.to_string().into()),
)
} else {
ServerJsonRpcMessage::error(
ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "Method not found", None),
Some(discover.id),
))
.await
.expect("send method-not-found");
)
};
server.send(reply).await.expect("send discovery reply");

let ClientJsonRpcMessage::Request(initialize) =
server.receive().await.expect("expected initialize request")
Expand All @@ -380,6 +411,15 @@ async fn auto_startup_falls_back_after_discover_method_not_found() {
server.receive().await,
Some(ClientJsonRpcMessage::Notification(_))
));
let (request, id) = server.receive().await.unwrap().into_request().unwrap();
assert!(matches!(request, ClientRequest::ListToolsRequest(_)));
server
.send(ServerJsonRpcMessage::response(
ServerResult::ListToolsResult(Default::default()),
id,
))
.await
.unwrap();
});

let client = DiscoverClient
Expand All @@ -392,6 +432,10 @@ async fn auto_startup_falls_back_after_discover_method_not_found() {
)
.await
.expect("auto client should fall back");
client
.list_tools(None)
.await
.expect("list tools after fallback");
client.cancel().await.expect("cancel client");
server_task.await.expect("server task");
}
Expand Down Expand Up @@ -588,8 +632,16 @@ async fn auto_startup_preserves_both_errors_when_fallback_also_fails() {
server_task.await.expect("server task");
}

#[tokio::test]
async fn auto_startup_surfaces_uncorrelated_error_id() {
#[rstest::rstest]
#[case::discovery(false)]
#[case::timeout_fallback(true)]
#[tokio::test(start_paused = true)]
async fn auto_startup_rejects_unrelated_responses(
#[case] after_timeout: bool,
#[values(false, true)] success: bool,
) {
use rmcp::service::ClientInitializeError;

let (server_transport, client_transport) = tokio::io::duplex(4096);
let mut server = IntoTransport::<rmcp::RoleServer, _, _>::into_transport(server_transport);
let server_task = tokio::spawn(async move {
Expand All @@ -598,14 +650,20 @@ async fn auto_startup_surfaces_uncorrelated_error_id() {
else {
panic!("expected request");
};
// Respond with an error carrying a different id.
server
.send(ServerJsonRpcMessage::error(
ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "not found", None),
Some(RequestId::Number(999)),
))
.await
.expect("send mismatched-id error");
if after_timeout {
let (request, _) = server.receive().await.unwrap().into_request().unwrap();
assert!(matches!(request, ClientRequest::InitializeRequest(_)));
}
let id = RequestId::Number(999);
let reply = if success {
ServerJsonRpcMessage::response(ServerResult::empty(()), id)
} else {
ServerJsonRpcMessage::error(
ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "Method not found", None),
Some(id),
)
};
server.send(reply).await.expect("send mismatched ID");
});

let result = DiscoverClient
Expand All @@ -618,14 +676,17 @@ async fn auto_startup_surfaces_uncorrelated_error_id() {
)
.await;

assert!(
matches!(
result,
Err(rmcp::service::ClientInitializeError::UncorrelatedErrorResponse { .. })
),
"mismatched id should surface as UncorrelatedErrorResponse, got {:?}",
result.as_ref().err()
);
let error = result
.err()
.expect("an unrelated response ID must fail startup");
assert!(matches!(
(success, error),
(true, ClientInitializeError::ConflictInitResponseId(..))
| (
false,
ClientInitializeError::UncorrelatedErrorResponse { .. }
)
));
server_task.await.expect("server task");
}

Expand Down
Loading