From e6ee5e097e2a19c7234eebdf405727bf9c1a99ca Mon Sep 17 00:00:00 2001 From: ip2a Date: Wed, 5 Aug 2026 15:39:56 +0800 Subject: [PATCH] fix(client): classify discover outcome at source, not at the error type MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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. --- crates/rmcp/src/service/client.rs | 174 +++++++--- .../rmcp/tests/test_client_initialization.rs | 11 +- .../rmcp/tests/test_client_lifecycle_modes.rs | 328 ++++++++++++++++++ 3 files changed, 468 insertions(+), 45 deletions(-) diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 2e25b13a3..54bfb9407 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -49,6 +49,14 @@ pub enum ClientInitializeError { #[error("conflict initialized response id: expected {0}, got {1}")] ConflictInitResponseId(RequestId, RequestId), + #[error( + "uncorrelated error response: expected id {expected}, error response carried {received}" + )] + UncorrelatedErrorResponse { + expected: RequestId, + received: RequestId, + }, + #[error("connection closed: {0}")] ConnectionClosed(String), @@ -74,6 +82,12 @@ pub enum ClientInitializeError { #[error("Cancelled")] Cancelled, + + #[error("discover and legacy initialize both failed")] + LegacyFallbackFailed { + discover: Box, + fallback: Box, + }, } impl ClientInitializeError { @@ -96,8 +110,13 @@ impl ClientInitializeError { pub fn auth_challenge(&self) -> Option<&str> { use crate::transport::streamable_http_client::{AuthRequiredError, InsufficientScopeError}; - let Self::TransportError { error, .. } = self else { - return None; + let error = match self { + Self::TransportError { error, .. } => error, + // A 401/403 in the fallback phase is still actionable. + Self::LegacyFallbackFailed { fallback, .. } => { + return fallback.auth_challenge(); + } + _ => return None, }; let mut source: Option<&(dyn std::error::Error + 'static)> = Some(error.error.as_ref()); while let Some(current) = source { @@ -117,10 +136,11 @@ impl ClientInitializeError { /// This covers both missing or expired local OAuth authorization and an HTTP /// authorization challenge from the MCP server. pub fn is_authorization_required(&self) -> bool { - matches!( - self, - Self::TransportError { error, .. } if error.is_authorization_required() - ) + match self { + Self::TransportError { error, .. } => error.is_authorization_required(), + Self::LegacyFallbackFailed { fallback, .. } => fallback.is_authorization_required(), + _ => false, + } } } @@ -138,13 +158,20 @@ where .ok_or_else(|| ClientInitializeError::ConnectionClosed(context.to_string())) } -/// Helper function to expect a response from the stream +/// Helper function to expect a response from the stream, correlated to +/// `expected_id`. +/// +/// 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. async fn expect_response( transport: &mut T, context: &str, service: &S, peer: Peer, -) -> Result<(ServerResult, RequestId), ClientInitializeError> + expected_id: &RequestId, +) -> Result where T: Transport, S: Service, @@ -152,13 +179,29 @@ where loop { let message = expect_next_message(transport, context).await?; match message { - // Expected message to complete the initialization ServerJsonRpcMessage::Response(JsonRpcResponse { id, result, .. }) => { - break Ok((result, id)); + if !expected_id.matches_response_id(&id) { + return Err(ClientInitializeError::ConflictInitResponseId( + expected_id.clone(), + id, + )); + } + return Ok(result); } - // Handle JSON-RPC error responses ServerJsonRpcMessage::Error(error) => { - break Err(ClientInitializeError::JsonRpcError(error.error)); + return Err(match &error.id { + Some(id) if expected_id.matches_response_id(id) => { + ClientInitializeError::JsonRpcError(error.error) + } + // Spec: error id is optional; a server that cannot read + // the request id omits it. The error is still a response + // to our request, so it remains available to the caller. + None => ClientInitializeError::JsonRpcError(error.error), + Some(id) => ClientInitializeError::UncorrelatedErrorResponse { + expected: expected_id.clone(), + received: id.clone(), + }, + }); } // Server could send logging messages before handshake ServerJsonRpcMessage::Notification(mut notification) => { @@ -714,7 +757,7 @@ where legacy_startup(&service, &mut transport, &id_provider, &peer, client_info).await?; } ClientLifecycleMode::Discover { preferred_versions } => { - discover_startup( + match discover_startup( &service, &mut transport, &id_provider, @@ -722,13 +765,18 @@ where &client_info, preferred_versions, ) - .await?; + .await? + { + DiscoverOutcome::Modern => {} + // Discover mode does not fall back; a legacy server is an error. + DiscoverOutcome::Legacy(error) => return Err(*error), + } } ClientLifecycleMode::Auto { preferred_versions, legacy_version, } => { - let discover_result = discover_startup( + match discover_startup( &service, &mut transport, &id_provider, @@ -736,18 +784,23 @@ where &client_info, preferred_versions, ) - .await; - match discover_result { - Ok(()) => {} - Err(ClientInitializeError::JsonRpcError(error)) - if error.code == crate::model::ErrorCode::METHOD_NOT_FOUND => - { + .await + { + Ok(DiscoverOutcome::Modern) => {} + Ok(DiscoverOutcome::Legacy(discover_error)) => { let mut legacy_info = client_info; if let Some(version) = legacy_version { legacy_info.protocol_version = version; } - legacy_startup(&service, &mut transport, &id_provider, &peer, legacy_info) - .await?; + if let Err(fallback_error) = + legacy_startup(&service, &mut transport, &id_provider, &peer, legacy_info) + .await + { + return Err(ClientInitializeError::LegacyFallbackFailed { + discover: discover_error, + fallback: Box::new(fallback_error), + }); + } } Err(error) => return Err(error), } @@ -756,6 +809,41 @@ where Ok(serve_inner(service, transport, peer, peer_rx, ct)) } +/// Modern-era JSON-RPC error codes a server can return from `server/discover` +/// without being legacy. Version negotiation (`UNSUPPORTED_PROTOCOL_VERSION`) +/// is handled by `discover_startup`'s own retry loop and never reaches the +/// classification below. +/// +/// `ErrorCode` is an open integer type, so this cannot be exhaustive: if a +/// future revision adds another modern-era rejection code, add it here. +fn is_modern_rejection_code(code: crate::model::ErrorCode) -> bool { + matches!( + code, + crate::model::ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY + | crate::model::ErrorCode::HEADER_MISMATCH + ) +} + +/// The outcome of a `server/discover` probe, classified at the point where all +/// the context (request id, response correlation, transport state) is still +/// available. +/// +/// `Legacy` is returned only when the probe produced a complete, correlated +/// JSON-RPC error whose code is not a modern-era rejection — i.e. the +/// transport is in a known-good state and the error identifies the peer as +/// legacy per the 2026-07-28 backward-compatibility guidance. Every other +/// failure (transport error, uncorrelated response, modern rejection, etc.) +/// becomes `Err` so the caller surfaces it instead of retrying. +enum DiscoverOutcome { + /// The server speaks the modern protocol; discovery succeeded. + Modern, + /// The server is legacy: discovery received a correlated, non-modern + /// JSON-RPC error. The transport is still usable for a legacy `initialize` + /// handshake. The original error is preserved so a failed fallback can + /// report both phases. + Legacy(Box), +} + async fn legacy_startup( service: &S, transport: &mut T, @@ -784,15 +872,8 @@ where context: "send initialize request".into(), })?; - let (response, response_id) = - expect_response(transport, "initialize response", service, peer.clone()).await?; - - if !id.matches_response_id(&response_id) { - return Err(ClientInitializeError::ConflictInitResponseId( - id, - response_id, - )); - } + let response = + expect_response(transport, "initialize response", service, peer.clone(), &id).await?; let ServerResult::InitializeResult(initialize_result) = response else { return Err(ClientInitializeError::ExpectedInitResult(Some(response))); @@ -819,7 +900,7 @@ async fn discover_startup( peer: &Peer, client_info: &ClientInfo, preferred_versions: Vec, -) -> Result<(), ClientInitializeError> +) -> Result where S: Service, T: Transport + 'static, @@ -851,14 +932,8 @@ where ClientInitializeError::transport::(error, "send discover request") })?; - match expect_response(transport, "discover response", service, peer.clone()).await { - Ok((ServerResult::DiscoverResult(result), response_id)) => { - if !id.matches_response_id(&response_id) { - return Err(ClientInitializeError::ConflictInitResponseId( - id, - response_id, - )); - } + match expect_response(transport, "discover response", service, peer.clone(), &id).await { + Ok(ServerResult::DiscoverResult(result)) => { let Some(selected) = select_protocol_version(&preferred_versions, &result.supported_versions) else { @@ -876,9 +951,9 @@ where client_info: client_info.client_info.clone(), client_capabilities: client_info.capabilities.clone(), }); - return Ok(()); + return Ok(DiscoverOutcome::Modern); } - Ok((response, _)) => { + Ok(response) => { return Err(ClientInitializeError::ExpectedInitResult(Some(response))); } Err(ClientInitializeError::JsonRpcError(error)) @@ -912,6 +987,19 @@ where }; candidate = next; } + // A correlated JSON-RPC error that is not a modern-era rejection + // and not a version-negotiation signal: the server is legacy. + // The transport delivered a complete response, so a legacy + // `initialize` can follow on the same connection. + Err(error) + if matches!( + &error, + ClientInitializeError::JsonRpcError(data) + if !is_modern_rejection_code(data.code) + ) => + { + return Ok(DiscoverOutcome::Legacy(Box::new(error))); + } Err(error) => return Err(error), } } diff --git a/crates/rmcp/tests/test_client_initialization.rs b/crates/rmcp/tests/test_client_initialization.rs index 960e1cf53..0ca0182f9 100644 --- a/crates/rmcp/tests/test_client_initialization.rs +++ b/crates/rmcp/tests/test_client_initialization.rs @@ -125,11 +125,18 @@ async fn test_client_init_handles_jsonrpc_error() { }); tokio::spawn(async move { - let _init_request = server.receive().await; + let request = server.receive().await; + // Echo the request's own id back on the error so it correlates: an + // uncorrelated id would surface as `UncorrelatedErrorResponse` + // instead of the `JsonRpcError` this test exercises. + let request_id = request + .and_then(|message| message.into_request()) + .map(|(_, id)| id) + .expect("client sent an initialize request"); let error_msg = ServerJsonRpcMessage::Error(JsonRpcError { jsonrpc: JsonRpcVersion2_0, - id: Some(RequestId::Number(1)), + id: Some(request_id), error: ErrorData { code: ErrorCode(-32600), message: Cow::Borrowed("Invalid Request"), diff --git a/crates/rmcp/tests/test_client_lifecycle_modes.rs b/crates/rmcp/tests/test_client_lifecycle_modes.rs index 4b75c3550..5fa2cd191 100644 --- a/crates/rmcp/tests/test_client_lifecycle_modes.rs +++ b/crates/rmcp/tests/test_client_lifecycle_modes.rs @@ -337,6 +337,334 @@ async fn auto_startup_falls_back_after_discover_method_not_found() { server_task.await.expect("server task"); } +/// Drives an `Auto` client through a single `server/discover` probe and asserts +/// the legacy fallback decision against the response the server sends back. +/// +/// When `expect_fallback` is set, the server also accepts the subsequent +/// `initialize` request and the client is expected to connect. Otherwise the +/// client must surface the discover error without sending `initialize`, and the +/// server's next receive must not be an initialize request. +async fn run_auto_discover_response_scenario(error: ErrorData, expect_fallback: 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 { + let ClientJsonRpcMessage::Request(discover) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + assert!(matches!( + discover.request, + ClientRequest::DiscoverRequest(_) + )); + server + .send(ServerJsonRpcMessage::error(error, Some(discover.id))) + .await + .expect("send discover error response"); + + if expect_fallback { + let ClientJsonRpcMessage::Request(initialize) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected request"); + }; + assert!(matches!( + initialize.request, + ClientRequest::InitializeRequest(_) + )); + server + .send(ServerJsonRpcMessage::response( + ServerResult::InitializeResult(InitializeResult::new( + ServerCapabilities::default(), + )), + initialize.id, + )) + .await + .expect("send initialize response"); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Notification(_)) + )); + } else { + // The client must surface the error without falling back, so no + // initialize request should follow. The transport closes when the + // failed client is dropped. + if let Some(ClientJsonRpcMessage::Request(request)) = server.receive().await { + panic!( + "client fell back to {:?} but should have surfaced the modern error", + request.request + ); + } + } + }); + + let client_result = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await; + + if expect_fallback { + let client = client_result.expect("auto client should fall back to initialize"); + client.cancel().await.expect("cancel client"); + } else { + assert!( + client_result.is_err(), + "modern error should surface without legacy fallback" + ); + } + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn auto_startup_falls_back_after_discover_invalid_request() { + // Legacy servers commonly reject an unknown pre-initialize request with + // `-32600` (e.g. a session middleware that requires `initialize` first). + run_auto_discover_response_scenario( + ErrorData::new(ErrorCode::INVALID_REQUEST, "Bad Request", None), + true, + ) + .await; +} + +#[tokio::test] +async fn auto_startup_falls_back_after_discover_invalid_params() { + // `-32602` is explicitly called out by the specification as an + // implementation-defined response legacy servers use for unknown requests. + run_auto_discover_response_scenario( + ErrorData::new(ErrorCode::INVALID_PARAMS, "Invalid params", None), + true, + ) + .await; +} + +#[tokio::test] +async fn auto_startup_does_not_fall_back_for_missing_required_capability() { + // A `MISSING_REQUIRED_CLIENT_CAPABILITY` response identifies a modern + // server; falling back to `initialize` would not address it. + run_auto_discover_response_scenario( + ErrorData::new( + ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY, + "Missing required client capability", + None, + ), + false, + ) + .await; +} + +#[tokio::test] +async fn auto_startup_does_not_fall_back_for_header_mismatch() { + // A `HEADER_MISMATCH` response identifies a modern server performing + // header validation; falling back to `initialize` would not address it. + run_auto_discover_response_scenario( + ErrorData::new(ErrorCode::HEADER_MISMATCH, "Header mismatch", None), + false, + ) + .await; +} + +#[tokio::test] +async fn auto_startup_preserves_both_errors_when_fallback_also_fails() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + // Discover: legacy rejection. + let ClientJsonRpcMessage::Request(discover) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + server + .send(ServerJsonRpcMessage::error( + ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "not found", None), + Some(discover.id), + )) + .await + .expect("send discover error"); + + // Initialize: close the transport instead of responding. + let ClientJsonRpcMessage::Request(_) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected request"); + }; + drop(server); // closes the transport + }); + + let result = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await; + + let err = result.err().expect("both phases should fail"); + match err { + rmcp::service::ClientInitializeError::LegacyFallbackFailed { discover, fallback } => { + assert!( + matches!( + *discover, + rmcp::service::ClientInitializeError::JsonRpcError(_) + ), + "discover phase should be a JsonRpcError, got {discover:?}" + ); + assert!( + matches!( + *fallback, + rmcp::service::ClientInitializeError::ConnectionClosed(_) + ), + "fallback phase should be ConnectionClosed, got {fallback:?}" + ); + } + other => panic!("expected LegacyFallbackFailed, got {other:?}"), + } + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn auto_startup_surfaces_uncorrelated_error_id() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(_) = + server.receive().await.expect("expected discover request") + 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"); + }); + + let result = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await; + + assert!( + matches!( + result, + Err(rmcp::service::ClientInitializeError::UncorrelatedErrorResponse { .. }) + ), + "mismatched id should surface as UncorrelatedErrorResponse, got {:?}", + result.as_ref().err() + ); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn auto_startup_falls_back_for_absent_error_id() { + // Spec: error id is optional; a server that cannot read the request id + // omits it. The error is still a response to our request, so it should + // trigger legacy fallback. + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(_discover) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + server + .send(ServerJsonRpcMessage::error( + ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "not found", None), + None, + )) + .await + .expect("send no-id error"); + + // Client should fall back to initialize. + let ClientJsonRpcMessage::Request(initialize) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected request"); + }; + server + .send(ServerJsonRpcMessage::response( + ServerResult::InitializeResult( + InitializeResult::new(ServerCapabilities::default()), + ), + initialize.id, + )) + .await + .expect("send initialize response"); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Notification(_)) + )); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Auto { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + }, + ) + .await + .expect("absent-id error should trigger legacy fallback"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn discover_mode_surfaces_legacy_error() { + let (server_transport, client_transport) = tokio::io::duplex(4096); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(discover) = + server.receive().await.expect("expected discover request") + else { + panic!("expected request"); + }; + server + .send(ServerJsonRpcMessage::error( + ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "not found", None), + Some(discover.id), + )) + .await + .expect("send discover error"); + }); + + let result = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await; + + assert!( + matches!( + result, + Err(rmcp::service::ClientInitializeError::JsonRpcError(_)) + ), + "Discover mode should surface a legacy error, not fall back, got {:?}", + result.as_ref().err() + ); + server_task.await.expect("server task"); +} + #[tokio::test] async fn discover_startup_retries_a_mutually_supported_version() { let unsupported: ProtocolVersion =