diff --git a/crates/rmcp/src/model.rs b/crates/rmcp/src/model.rs index 97273d9ae..09d296e9a 100644 --- a/crates/rmcp/src/model.rs +++ b/crates/rmcp/src/model.rs @@ -240,6 +240,22 @@ impl NumberOrString { NumberOrString::String(s) => Value::String(s.to_string()), } } + + pub(crate) fn numeric_string_value(&self) -> Option { + match self { + Self::String(id) => id.parse().ok(), + Self::Number(_) => None, + } + } + + pub(crate) fn matches_response_id(&self, response_id: &Self) -> bool { + self == response_id + || matches!( + self, + Self::Number(request_id) + if response_id.numeric_string_value() == Some(*request_id) + ) + } } impl std::fmt::Display for NumberOrString { diff --git a/crates/rmcp/src/service.rs b/crates/rmcp/src/service.rs index 7ef938e6d..4eeaeadf4 100644 --- a/crates/rmcp/src/service.rs +++ b/crates/rmcp/src/service.rs @@ -305,6 +305,17 @@ pub trait ProgressTokenProvider: Send + Sync + 'static { pub type AtomicU32RequestIdProvider = AtomicU32Provider; pub type AtomicU32ProgressTokenProvider = AtomicU32Provider; +pub(crate) fn remove_pending_request( + pending_requests: &mut HashMap, + response_id: &RequestId, +) -> Option { + pending_requests.remove(response_id).or_else(|| { + response_id + .numeric_string_value() + .and_then(|id| pending_requests.remove(&RequestId::Number(id))) + }) +} + #[derive(Debug, Default)] pub struct AtomicU32Provider { id: AtomicU64, @@ -1481,7 +1492,9 @@ where id, .. })) => { - if let Some(responder) = local_responder_pool.remove(&id) { + if let Some(responder) = + remove_pending_request(&mut local_responder_pool, &id) + { let response_result = responder.send(Ok(result)); if let Err(_error) = response_result { tracing::warn!(%id, "Error sending response"); @@ -1495,7 +1508,9 @@ where tracing::debug!(?error, "received id-less peer error"); continue; }; - if let Some(responder) = local_responder_pool.remove(&id) { + if let Some(responder) = + remove_pending_request(&mut local_responder_pool, &id) + { let service_error = if error.is_transport_closed() { ServiceError::TransportClosed } else { diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index 35deb38c8..093a537b1 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -687,7 +687,7 @@ where let (response, response_id) = expect_response(transport, "initialize response", service, peer.clone()).await?; - if id != response_id { + if !id.matches_response_id(&response_id) { return Err(ClientInitializeError::ConflictInitResponseId( id, response_id, @@ -753,7 +753,7 @@ where match expect_response(transport, "discover response", service, peer.clone()).await { Ok((ServerResult::DiscoverResult(result), response_id)) => { - if response_id != id { + if !id.matches_response_id(&response_id) { return Err(ClientInitializeError::ConflictInitResponseId( id, response_id, diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index 8c3a1f83d..432ff705c 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -499,8 +499,14 @@ impl StreamableHttpClientWorker { pending_stream_response_ids: &mut HashSet, message: &ServerJsonRpcMessage, ) { - if let Some(id) = Self::server_response_id(message) { - pending_stream_response_ids.remove(id); + let Some(response_id) = Self::server_response_id(message) else { + return; + }; + if pending_stream_response_ids.remove(response_id) { + return; + } + if let Some(id) = response_id.numeric_string_value() { + pending_stream_response_ids.remove(&RequestId::Number(id)); } } @@ -1382,7 +1388,10 @@ impl Worker for StreamableHttpClientWorker { } Event::ServerMessage(mut json_rpc_message) => { if let Some(response_id) = Self::server_response_id(&json_rpc_message) - && let Some(stream_ct) = request_stream_cancellations.remove(response_id) + && let Some(stream_ct) = crate::service::remove_pending_request( + &mut request_stream_cancellations, + response_id, + ) { stream_ct.cancel(); } @@ -1848,4 +1857,39 @@ mod tests { vec!["legacy"] ); } + + #[cfg(feature = "transport-streamable-http-client-reqwest")] + #[test] + fn clear_stream_response_pending_accepts_stringified_numeric_id() { + let mut pending = HashSet::from([NumberOrString::Number(1)]); + let response = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::default()), + NumberOrString::String("1".into()), + ); + + StreamableHttpClientWorker::::clear_stream_response_pending( + &mut pending, + &response, + ); + + assert!(pending.is_empty()); + } + + #[cfg(feature = "transport-streamable-http-client-reqwest")] + #[test] + fn clear_stream_response_pending_prefers_exact_string_id() { + let string_id = NumberOrString::String("1".into()); + let mut pending = HashSet::from([NumberOrString::Number(1), string_id.clone()]); + let response = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::default()), + string_id, + ); + + StreamableHttpClientWorker::::clear_stream_response_pending( + &mut pending, + &response, + ); + + assert_eq!(pending, HashSet::from([NumberOrString::Number(1)])); + } } diff --git a/crates/rmcp/tests/test_client_initialization.rs b/crates/rmcp/tests/test_client_initialization.rs index f51b33ef7..6c7984c5d 100644 --- a/crates/rmcp/tests/test_client_initialization.rs +++ b/crates/rmcp/tests/test_client_initialization.rs @@ -9,11 +9,102 @@ use common::handlers::TestClientHandler; use rmcp::{ ServiceExt, model::{ - ErrorCode, ErrorData, JsonRpcError, JsonRpcVersion2_0, RequestId, ServerJsonRpcMessage, + ClientJsonRpcMessage, ErrorCode, ErrorData, InitializeResult, JsonRpcError, + JsonRpcVersion2_0, RequestId, ServerCapabilities, ServerJsonRpcMessage, ServerResult, }, transport::{IntoTransport, Transport}, }; +fn stringify_numeric_id(id: RequestId) -> RequestId { + let RequestId::Number(id) = id else { + panic!("expected a numeric request ID"); + }; + RequestId::String(id.to_string().into()) +} + +#[tokio::test] +async fn client_initialization_accepts_stringified_numeric_response_id() { + let (server_transport, client_transport) = tokio::io::duplex(1024); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(request) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected initialize request"); + }; + server + .send(ServerJsonRpcMessage::response( + ServerResult::InitializeResult( + InitializeResult::new(ServerCapabilities::default()), + ), + stringify_numeric_id(request.id), + )) + .await + .expect("send initialize response"); + assert!(matches!( + server.receive().await, + Some(ClientJsonRpcMessage::Notification(_)) + )); + }); + + let client = TestClientHandler::new(true, true) + .serve(client_transport) + .await + .expect("client should accept stringified initialize response ID"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + +#[tokio::test] +async fn client_correlates_stringified_numeric_response_id() { + let (server_transport, client_transport) = tokio::io::duplex(1024); + let mut server = IntoTransport::::into_transport(server_transport); + let server_task = tokio::spawn(async move { + let ClientJsonRpcMessage::Request(initialize) = + server.receive().await.expect("expected initialize request") + else { + panic!("expected initialize 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 ClientJsonRpcMessage::Request(request) = + server.receive().await.expect("expected tools/list request") + else { + panic!("expected tools/list request"); + }; + server + .send(ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(Default::default()), + stringify_numeric_id(request.id), + )) + .await + .expect("send tools/list response"); + }); + + let client = TestClientHandler::new(true, true) + .serve(client_transport) + .await + .expect("initialize client"); + client + .list_tools(None) + .await + .expect("client should correlate stringified response ID"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + #[tokio::test] async fn test_client_init_handles_jsonrpc_error() { let (server_transport, client_transport) = tokio::io::duplex(1024); diff --git a/crates/rmcp/tests/test_client_lifecycle_modes.rs b/crates/rmcp/tests/test_client_lifecycle_modes.rs index 66638c1f2..375364e88 100644 --- a/crates/rmcp/tests/test_client_lifecycle_modes.rs +++ b/crates/rmcp/tests/test_client_lifecycle_modes.rs @@ -4,7 +4,7 @@ use rmcp::{ ClientHandler, ClientLifecycleMode, ClientServiceExt, ServerHandler, ServiceExt, model::{ ClientJsonRpcMessage, ClientRequest, DiscoverResult, ErrorCode, ErrorData, GetMeta, - Implementation, InitializeResult, ProtocolVersion, ServerCapabilities, + Implementation, InitializeResult, ProtocolVersion, RequestId, ServerCapabilities, ServerJsonRpcMessage, ServerResult, }, service::PeerRequestOptions, @@ -21,6 +21,45 @@ struct StatelessServer; impl ServerHandler for StatelessServer {} +#[tokio::test] +async fn discover_startup_accepts_stringified_numeric_response_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(request) = + server.receive().await.expect("expected discover request") + else { + panic!("expected discover request"); + }; + let RequestId::Number(response_id) = request.id else { + panic!("expected a numeric request ID"); + }; + server + .send(ServerJsonRpcMessage::response( + ServerResult::DiscoverResult(DiscoverResult::new( + vec![ProtocolVersion::V_2026_07_28], + ServerCapabilities::default(), + Implementation::new("discover-server", "1.0.0"), + )), + RequestId::String(response_id.to_string().into()), + )) + .await + .expect("send discover response"); + }); + + let client = DiscoverClient + .serve_with_lifecycle( + client_transport, + ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + }, + ) + .await + .expect("client should accept stringified discover response ID"); + client.cancel().await.expect("cancel client"); + server_task.await.expect("server task"); +} + #[tokio::test] async fn high_level_server_accepts_discover_startup_without_initialize() { let (server_transport, client_transport) = tokio::io::duplex(4096);