Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
16 changes: 16 additions & 0 deletions crates/rmcp/src/model.rs
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,22 @@ impl NumberOrString {
NumberOrString::String(s) => Value::String(s.to_string()),
}
}

pub(crate) fn numeric_string_value(&self) -> Option<i64> {
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 {
Expand Down
19 changes: 17 additions & 2 deletions crates/rmcp/src/service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -305,6 +305,17 @@ pub trait ProgressTokenProvider: Send + Sync + 'static {
pub type AtomicU32RequestIdProvider = AtomicU32Provider;
pub type AtomicU32ProgressTokenProvider = AtomicU32Provider;

pub(crate) fn remove_pending_request<T>(
pending_requests: &mut HashMap<RequestId, T>,
response_id: &RequestId,
) -> Option<T> {
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,
Expand Down Expand Up @@ -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");
Expand All @@ -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 {
Expand Down
4 changes: 2 additions & 2 deletions crates/rmcp/src/service/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
50 changes: 47 additions & 3 deletions crates/rmcp/src/transport/streamable_http_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -499,8 +499,14 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
pending_stream_response_ids: &mut HashSet<RequestId>,
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));
}
}

Expand Down Expand Up @@ -1382,7 +1388,10 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
}
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();
}
Expand Down Expand Up @@ -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::<reqwest::Client>::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::<reqwest::Client>::clear_stream_response_pending(
&mut pending,
&response,
);

assert_eq!(pending, HashSet::from([NumberOrString::Number(1)]));
}
}
93 changes: 92 additions & 1 deletion crates/rmcp/tests/test_client_initialization.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::<rmcp::RoleServer, _, _>::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::<rmcp::RoleServer, _, _>::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);
Expand Down
41 changes: 40 additions & 1 deletion crates/rmcp/tests/test_client_lifecycle_modes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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::<rmcp::RoleServer, _, _>::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);
Expand Down