Skip to content

Commit 22be2cc

Browse files
authored
fix(client): classify discover outcome at source, not at the error type (#1133)
`ClientLifecycleMode::Auto` only fell back from `server/discover` on `-32601`, so legacy servers that reject the probe with other codes (`-32600`, `-32602`, implementation-defined errors) failed to connect even though `initialize` would have succeeded. The previous attempt (indicates_legacy_server) classified the failure after the fact by reverse-engineering the error type. This rewrite moves the classification into `discover_startup` itself, where the full context (request id, response correlation, transport state) is still available. `discover_startup` now returns `DiscoverOutcome`: `Modern` on success, `Legacy(error)` when the probe received a complete, correlated JSON-RPC error whose code is not a modern-era rejection. Every other failure becomes `Err`, so `Auto` simply matches the outcome — no methods on `ClientInitializeError`, no downcast, no transport-specific types leaking into the generic lifecycle layer. Additional fixes that fall out naturally: - Response correlation is now checked in `expect_response` for both success and error branches. Previously error responses skipped id correlation entirely. A new `UncorrelatedErrorResponse` variant surfaces responses that cannot be tied to the request. - When both discover and the legacy fallback fail, a `LegacyFallbackFailed` compound error preserves both phases instead of discarding the discover error. Fixes #1040.
1 parent 02c62ae commit 22be2cc

3 files changed

Lines changed: 468 additions & 45 deletions

File tree

crates/rmcp/src/service/client.rs

Lines changed: 131 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,14 @@ pub enum ClientInitializeError {
4949
#[error("conflict initialized response id: expected {0}, got {1}")]
5050
ConflictInitResponseId(RequestId, RequestId),
5151

52+
#[error(
53+
"uncorrelated error response: expected id {expected}, error response carried {received}"
54+
)]
55+
UncorrelatedErrorResponse {
56+
expected: RequestId,
57+
received: RequestId,
58+
},
59+
5260
#[error("connection closed: {0}")]
5361
ConnectionClosed(String),
5462

@@ -74,6 +82,12 @@ pub enum ClientInitializeError {
7482

7583
#[error("Cancelled")]
7684
Cancelled,
85+
86+
#[error("discover and legacy initialize both failed")]
87+
LegacyFallbackFailed {
88+
discover: Box<ClientInitializeError>,
89+
fallback: Box<ClientInitializeError>,
90+
},
7791
}
7892

7993
impl ClientInitializeError {
@@ -96,8 +110,13 @@ impl ClientInitializeError {
96110
pub fn auth_challenge(&self) -> Option<&str> {
97111
use crate::transport::streamable_http_client::{AuthRequiredError, InsufficientScopeError};
98112

99-
let Self::TransportError { error, .. } = self else {
100-
return None;
113+
let error = match self {
114+
Self::TransportError { error, .. } => error,
115+
// A 401/403 in the fallback phase is still actionable.
116+
Self::LegacyFallbackFailed { fallback, .. } => {
117+
return fallback.auth_challenge();
118+
}
119+
_ => return None,
101120
};
102121
let mut source: Option<&(dyn std::error::Error + 'static)> = Some(error.error.as_ref());
103122
while let Some(current) = source {
@@ -117,10 +136,11 @@ impl ClientInitializeError {
117136
/// This covers both missing or expired local OAuth authorization and an HTTP
118137
/// authorization challenge from the MCP server.
119138
pub fn is_authorization_required(&self) -> bool {
120-
matches!(
121-
self,
122-
Self::TransportError { error, .. } if error.is_authorization_required()
123-
)
139+
match self {
140+
Self::TransportError { error, .. } => error.is_authorization_required(),
141+
Self::LegacyFallbackFailed { fallback, .. } => fallback.is_authorization_required(),
142+
_ => false,
143+
}
124144
}
125145
}
126146

@@ -138,27 +158,50 @@ where
138158
.ok_or_else(|| ClientInitializeError::ConnectionClosed(context.to_string()))
139159
}
140160

141-
/// Helper function to expect a response from the stream
161+
/// Helper function to expect a response from the stream, correlated to
162+
/// `expected_id`.
163+
///
164+
/// Both success and error responses are checked here: a mismatched id on a
165+
/// success response is `ConflictInitResponseId`; on an error response (whose
166+
/// `id` is optional per spec) it is `UncorrelatedErrorResponse`. The caller
167+
/// never sees an uncorrelated response.
142168
async fn expect_response<T, S>(
143169
transport: &mut T,
144170
context: &str,
145171
service: &S,
146172
peer: Peer<RoleClient>,
147-
) -> Result<(ServerResult, RequestId), ClientInitializeError>
173+
expected_id: &RequestId,
174+
) -> Result<ServerResult, ClientInitializeError>
148175
where
149176
T: Transport<RoleClient>,
150177
S: Service<RoleClient>,
151178
{
152179
loop {
153180
let message = expect_next_message(transport, context).await?;
154181
match message {
155-
// Expected message to complete the initialization
156182
ServerJsonRpcMessage::Response(JsonRpcResponse { id, result, .. }) => {
157-
break Ok((result, id));
183+
if !expected_id.matches_response_id(&id) {
184+
return Err(ClientInitializeError::ConflictInitResponseId(
185+
expected_id.clone(),
186+
id,
187+
));
188+
}
189+
return Ok(result);
158190
}
159-
// Handle JSON-RPC error responses
160191
ServerJsonRpcMessage::Error(error) => {
161-
break Err(ClientInitializeError::JsonRpcError(error.error));
192+
return Err(match &error.id {
193+
Some(id) if expected_id.matches_response_id(id) => {
194+
ClientInitializeError::JsonRpcError(error.error)
195+
}
196+
// Spec: error id is optional; a server that cannot read
197+
// the request id omits it. The error is still a response
198+
// to our request, so it remains available to the caller.
199+
None => ClientInitializeError::JsonRpcError(error.error),
200+
Some(id) => ClientInitializeError::UncorrelatedErrorResponse {
201+
expected: expected_id.clone(),
202+
received: id.clone(),
203+
},
204+
});
162205
}
163206
// Server could send logging messages before handshake
164207
ServerJsonRpcMessage::Notification(mut notification) => {
@@ -714,40 +757,50 @@ where
714757
legacy_startup(&service, &mut transport, &id_provider, &peer, client_info).await?;
715758
}
716759
ClientLifecycleMode::Discover { preferred_versions } => {
717-
discover_startup(
760+
match discover_startup(
718761
&service,
719762
&mut transport,
720763
&id_provider,
721764
&peer,
722765
&client_info,
723766
preferred_versions,
724767
)
725-
.await?;
768+
.await?
769+
{
770+
DiscoverOutcome::Modern => {}
771+
// Discover mode does not fall back; a legacy server is an error.
772+
DiscoverOutcome::Legacy(error) => return Err(*error),
773+
}
726774
}
727775
ClientLifecycleMode::Auto {
728776
preferred_versions,
729777
legacy_version,
730778
} => {
731-
let discover_result = discover_startup(
779+
match discover_startup(
732780
&service,
733781
&mut transport,
734782
&id_provider,
735783
&peer,
736784
&client_info,
737785
preferred_versions,
738786
)
739-
.await;
740-
match discover_result {
741-
Ok(()) => {}
742-
Err(ClientInitializeError::JsonRpcError(error))
743-
if error.code == crate::model::ErrorCode::METHOD_NOT_FOUND =>
744-
{
787+
.await
788+
{
789+
Ok(DiscoverOutcome::Modern) => {}
790+
Ok(DiscoverOutcome::Legacy(discover_error)) => {
745791
let mut legacy_info = client_info;
746792
if let Some(version) = legacy_version {
747793
legacy_info.protocol_version = version;
748794
}
749-
legacy_startup(&service, &mut transport, &id_provider, &peer, legacy_info)
750-
.await?;
795+
if let Err(fallback_error) =
796+
legacy_startup(&service, &mut transport, &id_provider, &peer, legacy_info)
797+
.await
798+
{
799+
return Err(ClientInitializeError::LegacyFallbackFailed {
800+
discover: discover_error,
801+
fallback: Box::new(fallback_error),
802+
});
803+
}
751804
}
752805
Err(error) => return Err(error),
753806
}
@@ -756,6 +809,41 @@ where
756809
Ok(serve_inner(service, transport, peer, peer_rx, ct))
757810
}
758811

812+
/// Modern-era JSON-RPC error codes a server can return from `server/discover`
813+
/// without being legacy. Version negotiation (`UNSUPPORTED_PROTOCOL_VERSION`)
814+
/// is handled by `discover_startup`'s own retry loop and never reaches the
815+
/// classification below.
816+
///
817+
/// `ErrorCode` is an open integer type, so this cannot be exhaustive: if a
818+
/// future revision adds another modern-era rejection code, add it here.
819+
fn is_modern_rejection_code(code: crate::model::ErrorCode) -> bool {
820+
matches!(
821+
code,
822+
crate::model::ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY
823+
| crate::model::ErrorCode::HEADER_MISMATCH
824+
)
825+
}
826+
827+
/// The outcome of a `server/discover` probe, classified at the point where all
828+
/// the context (request id, response correlation, transport state) is still
829+
/// available.
830+
///
831+
/// `Legacy` is returned only when the probe produced a complete, correlated
832+
/// JSON-RPC error whose code is not a modern-era rejection — i.e. the
833+
/// transport is in a known-good state and the error identifies the peer as
834+
/// legacy per the 2026-07-28 backward-compatibility guidance. Every other
835+
/// failure (transport error, uncorrelated response, modern rejection, etc.)
836+
/// becomes `Err` so the caller surfaces it instead of retrying.
837+
enum DiscoverOutcome {
838+
/// The server speaks the modern protocol; discovery succeeded.
839+
Modern,
840+
/// The server is legacy: discovery received a correlated, non-modern
841+
/// JSON-RPC error. The transport is still usable for a legacy `initialize`
842+
/// handshake. The original error is preserved so a failed fallback can
843+
/// report both phases.
844+
Legacy(Box<ClientInitializeError>),
845+
}
846+
759847
async fn legacy_startup<S, T>(
760848
service: &S,
761849
transport: &mut T,
@@ -784,15 +872,8 @@ where
784872
context: "send initialize request".into(),
785873
})?;
786874

787-
let (response, response_id) =
788-
expect_response(transport, "initialize response", service, peer.clone()).await?;
789-
790-
if !id.matches_response_id(&response_id) {
791-
return Err(ClientInitializeError::ConflictInitResponseId(
792-
id,
793-
response_id,
794-
));
795-
}
875+
let response =
876+
expect_response(transport, "initialize response", service, peer.clone(), &id).await?;
796877

797878
let ServerResult::InitializeResult(initialize_result) = response else {
798879
return Err(ClientInitializeError::ExpectedInitResult(Some(response)));
@@ -819,7 +900,7 @@ async fn discover_startup<S, T>(
819900
peer: &Peer<RoleClient>,
820901
client_info: &ClientInfo,
821902
preferred_versions: Vec<ProtocolVersion>,
822-
) -> Result<(), ClientInitializeError>
903+
) -> Result<DiscoverOutcome, ClientInitializeError>
823904
where
824905
S: Service<RoleClient>,
825906
T: Transport<RoleClient> + 'static,
@@ -851,14 +932,8 @@ where
851932
ClientInitializeError::transport::<T>(error, "send discover request")
852933
})?;
853934

854-
match expect_response(transport, "discover response", service, peer.clone()).await {
855-
Ok((ServerResult::DiscoverResult(result), response_id)) => {
856-
if !id.matches_response_id(&response_id) {
857-
return Err(ClientInitializeError::ConflictInitResponseId(
858-
id,
859-
response_id,
860-
));
861-
}
935+
match expect_response(transport, "discover response", service, peer.clone(), &id).await {
936+
Ok(ServerResult::DiscoverResult(result)) => {
862937
let Some(selected) =
863938
select_protocol_version(&preferred_versions, &result.supported_versions)
864939
else {
@@ -876,9 +951,9 @@ where
876951
client_info: client_info.client_info.clone(),
877952
client_capabilities: client_info.capabilities.clone(),
878953
});
879-
return Ok(());
954+
return Ok(DiscoverOutcome::Modern);
880955
}
881-
Ok((response, _)) => {
956+
Ok(response) => {
882957
return Err(ClientInitializeError::ExpectedInitResult(Some(response)));
883958
}
884959
Err(ClientInitializeError::JsonRpcError(error))
@@ -912,6 +987,19 @@ where
912987
};
913988
candidate = next;
914989
}
990+
// A correlated JSON-RPC error that is not a modern-era rejection
991+
// and not a version-negotiation signal: the server is legacy.
992+
// The transport delivered a complete response, so a legacy
993+
// `initialize` can follow on the same connection.
994+
Err(error)
995+
if matches!(
996+
&error,
997+
ClientInitializeError::JsonRpcError(data)
998+
if !is_modern_rejection_code(data.code)
999+
) =>
1000+
{
1001+
return Ok(DiscoverOutcome::Legacy(Box::new(error)));
1002+
}
9151003
Err(error) => return Err(error),
9161004
}
9171005
}

crates/rmcp/tests/test_client_initialization.rs

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,11 +125,18 @@ async fn test_client_init_handles_jsonrpc_error() {
125125
});
126126

127127
tokio::spawn(async move {
128-
let _init_request = server.receive().await;
128+
let request = server.receive().await;
129+
// Echo the request's own id back on the error so it correlates: an
130+
// uncorrelated id would surface as `UncorrelatedErrorResponse`
131+
// instead of the `JsonRpcError` this test exercises.
132+
let request_id = request
133+
.and_then(|message| message.into_request())
134+
.map(|(_, id)| id)
135+
.expect("client sent an initialize request");
129136

130137
let error_msg = ServerJsonRpcMessage::Error(JsonRpcError {
131138
jsonrpc: JsonRpcVersion2_0,
132-
id: Some(RequestId::Number(1)),
139+
id: Some(request_id),
133140
error: ErrorData {
134141
code: ErrorCode(-32600),
135142
message: Cow::Borrowed("Invalid Request"),

0 commit comments

Comments
 (0)