fix: run first pre-init request in service loop (#1263)
This commit is contained in:
parent
47fcdf4b57
commit
d481cb0eec
4 changed files with 121 additions and 23 deletions
|
|
@ -1298,7 +1298,14 @@ where
|
||||||
{
|
{
|
||||||
let (peer, peer_rx) = Peer::new(Arc::new(AtomicU32RequestIdProvider::default()), peer_info);
|
let (peer, peer_rx) = Peer::new(Arc::new(AtomicU32RequestIdProvider::default()), peer_info);
|
||||||
R::configure_direct_peer(&peer, &service.get_info());
|
R::configure_direct_peer(&peer, &service.get_info());
|
||||||
serve_inner(service, transport.into_transport(), peer, peer_rx, ct)
|
serve_inner(
|
||||||
|
service,
|
||||||
|
transport.into_transport(),
|
||||||
|
peer,
|
||||||
|
peer_rx,
|
||||||
|
VecDeque::new(),
|
||||||
|
ct,
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Spawn a task that may hold `!Send` state when the `local` feature is active.
|
/// Spawn a task that may hold `!Send` state when the `local` feature is active.
|
||||||
|
|
@ -1323,12 +1330,19 @@ where
|
||||||
tokio::task::spawn_local(future)
|
tokio::task::spawn_local(future)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Run the service loop over `transport`.
|
||||||
|
///
|
||||||
|
/// `initial_messages` are messages the caller already read from `transport`
|
||||||
|
/// (e.g. the first request of an `initialize`-less session). They are
|
||||||
|
/// dispatched by the loop, in order, before anything else is read from the
|
||||||
|
/// transport, so that their handlers run with the loop draining `peer_rx`.
|
||||||
#[instrument(skip_all)]
|
#[instrument(skip_all)]
|
||||||
fn serve_inner<R, S, T>(
|
fn serve_inner<R, S, T>(
|
||||||
service: S,
|
service: S,
|
||||||
transport: T,
|
transport: T,
|
||||||
peer: Peer<R>,
|
peer: Peer<R>,
|
||||||
mut peer_rx: tokio::sync::mpsc::Receiver<PeerSinkMessage<R>>,
|
mut peer_rx: tokio::sync::mpsc::Receiver<PeerSinkMessage<R>>,
|
||||||
|
initial_messages: VecDeque<RxJsonRpcMessage<R>>,
|
||||||
ct: CancellationToken,
|
ct: CancellationToken,
|
||||||
) -> RunningService<R, S>
|
) -> RunningService<R, S>
|
||||||
where
|
where
|
||||||
|
|
@ -1361,7 +1375,7 @@ where
|
||||||
let current_span = tracing::Span::current();
|
let current_span = tracing::Span::current();
|
||||||
let handle = spawn_service_task(async move {
|
let handle = spawn_service_task(async move {
|
||||||
let mut transport = transport.into_transport();
|
let mut transport = transport.into_transport();
|
||||||
let mut batch_messages = VecDeque::<RxJsonRpcMessage<R>>::new();
|
let mut batch_messages = initial_messages;
|
||||||
let mut send_task_set = tokio::task::JoinSet::<SendTaskResult>::new();
|
let mut send_task_set = tokio::task::JoinSet::<SendTaskResult>::new();
|
||||||
let mut response_send_tasks = tokio::task::JoinSet::<()>::new();
|
let mut response_send_tasks = tokio::task::JoinSet::<()>::new();
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
|
|
|
||||||
|
|
@ -827,7 +827,14 @@ where
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Ok(serve_inner(service, transport, peer, peer_rx, ct))
|
Ok(serve_inner(
|
||||||
|
service,
|
||||||
|
transport,
|
||||||
|
peer,
|
||||||
|
peer_rx,
|
||||||
|
VecDeque::new(),
|
||||||
|
ct,
|
||||||
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Modern-era JSON-RPC error codes a server can return from `server/discover`
|
/// Modern-era JSON-RPC error codes a server can return from `server/discover`
|
||||||
|
|
|
||||||
|
|
@ -594,7 +594,7 @@ where
|
||||||
|
|
||||||
let initialize_request = match request {
|
let initialize_request = match request {
|
||||||
ClientRequest::InitializeRequest(request) => request,
|
ClientRequest::InitializeRequest(request) => request,
|
||||||
mut request => {
|
request => {
|
||||||
let missing_metadata = request
|
let missing_metadata = request
|
||||||
.get_meta()
|
.get_meta()
|
||||||
.missing_required_keys(&ProtocolVersion::V_2026_07_28);
|
.missing_required_keys(&ProtocolVersion::V_2026_07_28);
|
||||||
|
|
@ -617,21 +617,17 @@ where
|
||||||
}
|
}
|
||||||
let (peer, peer_rx) = Peer::new(id_provider, None);
|
let (peer, peer_rx) = Peer::new(id_provider, None);
|
||||||
peer.require_request_metadata();
|
peer.require_request_metadata();
|
||||||
let context = RequestContext {
|
// Dispatch the request from inside the service loop rather than
|
||||||
ct: ct.child_token(),
|
// inline: its handler may send notifications through `peer`, which
|
||||||
id: id.clone(),
|
// only complete once the loop drains `peer_rx`.
|
||||||
meta: std::mem::take(request.get_meta_mut()),
|
return Ok(serve_inner(
|
||||||
extensions: std::mem::take(request.extensions_mut()),
|
service,
|
||||||
peer: peer.clone(),
|
transport,
|
||||||
};
|
peer,
|
||||||
let response = match service.handle_request(request, context).await {
|
peer_rx,
|
||||||
Ok(result) => ServerJsonRpcMessage::response(result, id),
|
VecDeque::from([ClientJsonRpcMessage::request(request, id)]),
|
||||||
Err(error) => ServerJsonRpcMessage::error(error, Some(id)),
|
ct,
|
||||||
};
|
));
|
||||||
transport.send(response).await.map_err(|error| {
|
|
||||||
ServerInitializeError::transport::<T>(error, "sending negotiated request response")
|
|
||||||
})?;
|
|
||||||
return Ok(serve_inner(service, transport, peer, peer_rx, ct));
|
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
let requested_protocol_version = initialize_request.params.protocol_version.clone();
|
let requested_protocol_version = initialize_request.params.protocol_version.clone();
|
||||||
|
|
@ -681,7 +677,14 @@ where
|
||||||
// Streamable HTTP has no ordering guarantee between POSTs, and the MCP spec uses
|
// Streamable HTTP has no ordering guarantee between POSTs, and the MCP spec uses
|
||||||
// SHOULD NOT (not MUST NOT) for pre-initialized messages, so any request arriving
|
// SHOULD NOT (not MUST NOT) for pre-initialized messages, so any request arriving
|
||||||
// before initialized is processed normally.
|
// before initialized is processed normally.
|
||||||
Ok(serve_inner(service, transport, peer, peer_rx, ct))
|
Ok(serve_inner(
|
||||||
|
service,
|
||||||
|
transport,
|
||||||
|
peer,
|
||||||
|
peer_rx,
|
||||||
|
VecDeque::new(),
|
||||||
|
ct,
|
||||||
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
macro_rules! method {
|
macro_rules! method {
|
||||||
|
|
|
||||||
|
|
@ -1,14 +1,18 @@
|
||||||
#![cfg(all(feature = "server", not(feature = "local")))]
|
#![cfg(all(feature = "server", not(feature = "local")))]
|
||||||
|
|
||||||
use std::sync::{Arc, Mutex};
|
use std::{
|
||||||
|
sync::{Arc, Mutex},
|
||||||
|
time::Duration,
|
||||||
|
};
|
||||||
|
|
||||||
use rmcp::{
|
use rmcp::{
|
||||||
ServerHandler, ServiceExt,
|
ServerHandler, ServiceExt,
|
||||||
model::{
|
model::{
|
||||||
ClientCapabilities, ClientJsonRpcMessage, ClientRequest, DiscoverRequest,
|
ClientCapabilities, ClientJsonRpcMessage, ClientRequest, DiscoverRequest,
|
||||||
DiscoverRequestParams, ErrorCode, ErrorData, Implementation, ListToolsRequest,
|
DiscoverRequestParams, ErrorCode, ErrorData, Implementation, ListToolsRequest,
|
||||||
ListToolsResult, PaginatedRequestParams, ProtocolVersion, RequestId, RequestMetaObject,
|
ListToolsResult, NumberOrString, PaginatedRequestParams, ProgressNotificationParam,
|
||||||
ServerJsonRpcMessage,
|
ProgressToken, ProtocolVersion, RequestId, RequestMetaObject, ServerJsonRpcMessage,
|
||||||
|
ServerNotification,
|
||||||
},
|
},
|
||||||
service::{MaybeSendFuture, RequestContext, RoleServer, ServerInitializeError},
|
service::{MaybeSendFuture, RequestContext, RoleServer, ServerInitializeError},
|
||||||
transport::{IntoTransport, Transport},
|
transport::{IntoTransport, Transport},
|
||||||
|
|
@ -212,3 +216,73 @@ async fn stateless_server_rejects_malformed_metadata_opener_with_error_response(
|
||||||
ServerInitializeError::ExpectedInitializeRequest(Some(_))
|
ServerInitializeError::ExpectedInitializeRequest(Some(_))
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct ProgressServer;
|
||||||
|
|
||||||
|
impl ServerHandler for ProgressServer {
|
||||||
|
async fn list_tools(
|
||||||
|
&self,
|
||||||
|
_request: Option<PaginatedRequestParams>,
|
||||||
|
context: RequestContext<RoleServer>,
|
||||||
|
) -> Result<ListToolsResult, ErrorData> {
|
||||||
|
let progress_token = context
|
||||||
|
.meta
|
||||||
|
.get_progress_token()
|
||||||
|
.expect("progress token in request meta");
|
||||||
|
context
|
||||||
|
.peer
|
||||||
|
.notify_progress(ProgressNotificationParam::new(progress_token, 1.0))
|
||||||
|
.await
|
||||||
|
.expect("send progress notification");
|
||||||
|
Ok(ListToolsResult::default())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Regression test for issue #1261: the first request of an `initialize`-less
|
||||||
|
/// session used to be handled outside the service loop, so a handler that sent
|
||||||
|
/// a notification before returning waited forever for the loop to flush it.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn stateless_server_first_request_handler_can_send_notifications() {
|
||||||
|
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
||||||
|
let server_task = tokio::spawn(async move {
|
||||||
|
ProgressServer
|
||||||
|
.serve(server_transport)
|
||||||
|
.await
|
||||||
|
.expect("server should start")
|
||||||
|
});
|
||||||
|
let mut client = IntoTransport::<rmcp::RoleClient, _, _>::into_transport(client_transport);
|
||||||
|
|
||||||
|
let mut meta = complete_meta();
|
||||||
|
meta.set_progress_token(ProgressToken(NumberOrString::Number(7)));
|
||||||
|
client
|
||||||
|
.send(list_tools_request(meta))
|
||||||
|
.await
|
||||||
|
.expect("send first request");
|
||||||
|
|
||||||
|
let exchange = async {
|
||||||
|
let notification = match client.receive().await {
|
||||||
|
Some(ServerJsonRpcMessage::Notification(notification)) => notification.notification,
|
||||||
|
other => panic!("expected progress notification before the response, got {other:?}"),
|
||||||
|
};
|
||||||
|
assert!(
|
||||||
|
matches!(notification, ServerNotification::ProgressNotification(_)),
|
||||||
|
"expected progress notification, got {notification:?}"
|
||||||
|
);
|
||||||
|
let response = match client.receive().await {
|
||||||
|
Some(ServerJsonRpcMessage::Response(response)) => response,
|
||||||
|
other => panic!("expected list tools response, got {other:?}"),
|
||||||
|
};
|
||||||
|
assert_eq!(response.id, RequestId::Number(1));
|
||||||
|
};
|
||||||
|
tokio::time::timeout(Duration::from_secs(5), exchange)
|
||||||
|
.await
|
||||||
|
.expect("first request must not deadlock the server");
|
||||||
|
|
||||||
|
server_task
|
||||||
|
.await
|
||||||
|
.expect("server task")
|
||||||
|
.cancel()
|
||||||
|
.await
|
||||||
|
.expect("cancel server");
|
||||||
|
}
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue