fix: accept stringified numeric response IDs (#1021)
This commit is contained in:
parent
99ee024499
commit
8803d39034
6 changed files with 214 additions and 9 deletions
|
|
@ -240,6 +240,22 @@ impl NumberOrString {
|
||||||
NumberOrString::String(s) => Value::String(s.to_string()),
|
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 {
|
impl std::fmt::Display for NumberOrString {
|
||||||
|
|
|
||||||
|
|
@ -305,6 +305,17 @@ pub trait ProgressTokenProvider: Send + Sync + 'static {
|
||||||
pub type AtomicU32RequestIdProvider = AtomicU32Provider;
|
pub type AtomicU32RequestIdProvider = AtomicU32Provider;
|
||||||
pub type AtomicU32ProgressTokenProvider = 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)]
|
#[derive(Debug, Default)]
|
||||||
pub struct AtomicU32Provider {
|
pub struct AtomicU32Provider {
|
||||||
id: AtomicU64,
|
id: AtomicU64,
|
||||||
|
|
@ -1481,7 +1492,9 @@ where
|
||||||
id,
|
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));
|
let response_result = responder.send(Ok(result));
|
||||||
if let Err(_error) = response_result {
|
if let Err(_error) = response_result {
|
||||||
tracing::warn!(%id, "Error sending response");
|
tracing::warn!(%id, "Error sending response");
|
||||||
|
|
@ -1495,7 +1508,9 @@ where
|
||||||
tracing::debug!(?error, "received id-less peer error");
|
tracing::debug!(?error, "received id-less peer error");
|
||||||
continue;
|
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() {
|
let service_error = if error.is_transport_closed() {
|
||||||
ServiceError::TransportClosed
|
ServiceError::TransportClosed
|
||||||
} else {
|
} else {
|
||||||
|
|
|
||||||
|
|
@ -687,7 +687,7 @@ where
|
||||||
let (response, response_id) =
|
let (response, response_id) =
|
||||||
expect_response(transport, "initialize response", service, peer.clone()).await?;
|
expect_response(transport, "initialize response", service, peer.clone()).await?;
|
||||||
|
|
||||||
if id != response_id {
|
if !id.matches_response_id(&response_id) {
|
||||||
return Err(ClientInitializeError::ConflictInitResponseId(
|
return Err(ClientInitializeError::ConflictInitResponseId(
|
||||||
id,
|
id,
|
||||||
response_id,
|
response_id,
|
||||||
|
|
@ -753,7 +753,7 @@ where
|
||||||
|
|
||||||
match expect_response(transport, "discover response", service, peer.clone()).await {
|
match expect_response(transport, "discover response", service, peer.clone()).await {
|
||||||
Ok((ServerResult::DiscoverResult(result), response_id)) => {
|
Ok((ServerResult::DiscoverResult(result), response_id)) => {
|
||||||
if response_id != id {
|
if !id.matches_response_id(&response_id) {
|
||||||
return Err(ClientInitializeError::ConflictInitResponseId(
|
return Err(ClientInitializeError::ConflictInitResponseId(
|
||||||
id,
|
id,
|
||||||
response_id,
|
response_id,
|
||||||
|
|
|
||||||
|
|
@ -499,8 +499,14 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
||||||
pending_stream_response_ids: &mut HashSet<RequestId>,
|
pending_stream_response_ids: &mut HashSet<RequestId>,
|
||||||
message: &ServerJsonRpcMessage,
|
message: &ServerJsonRpcMessage,
|
||||||
) {
|
) {
|
||||||
if let Some(id) = Self::server_response_id(message) {
|
let Some(response_id) = Self::server_response_id(message) else {
|
||||||
pending_stream_response_ids.remove(id);
|
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<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
||||||
}
|
}
|
||||||
Event::ServerMessage(mut json_rpc_message) => {
|
Event::ServerMessage(mut json_rpc_message) => {
|
||||||
if let Some(response_id) = Self::server_response_id(&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();
|
stream_ct.cancel();
|
||||||
}
|
}
|
||||||
|
|
@ -1848,4 +1857,39 @@ mod tests {
|
||||||
vec!["legacy"]
|
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)]));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -9,11 +9,102 @@ use common::handlers::TestClientHandler;
|
||||||
use rmcp::{
|
use rmcp::{
|
||||||
ServiceExt,
|
ServiceExt,
|
||||||
model::{
|
model::{
|
||||||
ErrorCode, ErrorData, JsonRpcError, JsonRpcVersion2_0, RequestId, ServerJsonRpcMessage,
|
ClientJsonRpcMessage, ErrorCode, ErrorData, InitializeResult, JsonRpcError,
|
||||||
|
JsonRpcVersion2_0, RequestId, ServerCapabilities, ServerJsonRpcMessage, ServerResult,
|
||||||
},
|
},
|
||||||
transport::{IntoTransport, Transport},
|
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]
|
#[tokio::test]
|
||||||
async fn test_client_init_handles_jsonrpc_error() {
|
async fn test_client_init_handles_jsonrpc_error() {
|
||||||
let (server_transport, client_transport) = tokio::io::duplex(1024);
|
let (server_transport, client_transport) = tokio::io::duplex(1024);
|
||||||
|
|
|
||||||
|
|
@ -4,7 +4,7 @@ use rmcp::{
|
||||||
ClientHandler, ClientLifecycleMode, ClientServiceExt, ServerHandler, ServiceExt,
|
ClientHandler, ClientLifecycleMode, ClientServiceExt, ServerHandler, ServiceExt,
|
||||||
model::{
|
model::{
|
||||||
ClientJsonRpcMessage, ClientRequest, DiscoverResult, ErrorCode, ErrorData, GetMeta,
|
ClientJsonRpcMessage, ClientRequest, DiscoverResult, ErrorCode, ErrorData, GetMeta,
|
||||||
Implementation, InitializeResult, ProtocolVersion, ServerCapabilities,
|
Implementation, InitializeResult, ProtocolVersion, RequestId, ServerCapabilities,
|
||||||
ServerJsonRpcMessage, ServerResult,
|
ServerJsonRpcMessage, ServerResult,
|
||||||
},
|
},
|
||||||
service::PeerRequestOptions,
|
service::PeerRequestOptions,
|
||||||
|
|
@ -21,6 +21,45 @@ struct StatelessServer;
|
||||||
|
|
||||||
impl ServerHandler for 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]
|
#[tokio::test]
|
||||||
async fn high_level_server_accepts_discover_startup_without_initialize() {
|
async fn high_level_server_accepts_discover_startup_without_initialize() {
|
||||||
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue