From c231e448136790cef2d7b945095bed2d89cd25d5 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 6 Nov 2026 14:16:51 -0500 Subject: [PATCH] fix(client): ignore late discovery replies after legacy fallback --- crates/rmcp/src/service/client.rs | 45 +++++-- .../rmcp/tests/test_client_lifecycle_modes.rs | 111 ++++++++++++++---- 2 files changed, 118 insertions(+), 38 deletions(-) diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 21d68be89..127bb9dd1 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -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: @@ -187,15 +188,21 @@ async fn expect_response( peer: Peer, expected_id: &RequestId, pre_init_ping: PreInitPing, + timed_out_probe: Option<&RequestId>, ) -> Result where T: Transport + 'static, S: Service, { + 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(), @@ -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) @@ -837,19 +847,19 @@ where T: Transport + 'static, { let mut transport = transport.into_transport(); - let id_provider = >::default(); - let (peer, peer_rx) = Peer::new(id_provider.clone(), None); + let (peer, peer_rx) = Peer::new(Arc::::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, @@ -870,7 +880,7 @@ where discover_startup( &service, &mut transport, - &id_provider, + &mut last_discover_id, &peer, &client_info, preferred_versions, @@ -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, @@ -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?; } } } @@ -954,15 +969,15 @@ enum DiscoverOutcome { async fn legacy_startup( service: &S, transport: &mut T, - id_provider: &Arc, peer: &Peer, client_info: InitializeRequestParams, + timed_out_probe: Option<&RequestId>, ) -> Result<(), ClientInitializeError> where S: Service, T: Transport + '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(), @@ -987,6 +1002,7 @@ where peer.clone(), &id, PreInitPing::Answer, + timed_out_probe, ) .await?; @@ -1022,7 +1038,7 @@ where async fn discover_startup( service: &S, transport: &mut T, - id_provider: &Arc, + last_discover_id: &mut Option, peer: &Peer, client_info: &InitializeRequestParams, preferred_versions: Vec, @@ -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 { @@ -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), @@ -1065,6 +1083,7 @@ where peer.clone(), &id, PreInitPing::Ignore, + None, ) .await { diff --git a/crates/rmcp/tests/test_client_lifecycle_modes.rs b/crates/rmcp/tests/test_client_lifecycle_modes.rs index 13ff77c8a..d4922ccb6 100644 --- a/crates/rmcp/tests/test_client_lifecycle_modes.rs +++ b/crates/rmcp/tests/test_client_lifecycle_modes.rs @@ -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::::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 { @@ -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") @@ -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 @@ -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"); } @@ -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::::into_transport(server_transport); let server_task = tokio::spawn(async move { @@ -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 @@ -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"); }