fix: accept logging/setLevel and ping before initialized notification (#730)

* fix: accept logging/setLevel and ping before initialized notification

* test: add server initialization tests for pre-init requests
This commit is contained in:
Dale Seo 2026-03-09 07:08:16 -04:00 • committed by GitHub
commit 8e5ebb4f5c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -16,7 +16,7 @@ use crate::{
model::{
CancelledNotification, CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage,
ClientNotification, ClientRequest, ClientResult, CreateMessageRequest,
CreateMessageRequestParams, CreateMessageResult, ErrorData, ListRootsRequest,
CreateMessageRequestParams, CreateMessageResult, EmptyResult, ErrorData, ListRootsRequest,
ListRootsResult, LoggingMessageNotification, LoggingMessageNotificationParam,
ProgressNotification, ProgressNotificationParam, PromptListChangedNotification,
ProtocolVersion, ResourceListChangedNotification, ResourceUpdatedNotification,
@ -147,22 +147,6 @@ where
)))
}
/// Helper function to expect a notification from the stream
async fn expect_notification<T>(
transport: &mut T,
context: &str,
) -> Result<ClientNotification, ServerInitializeError>
where
T: Transport<RoleServer>,
{
let msg = expect_next_message(transport, context).await?;
let msg_clone = msg.clone();
msg.into_notification()
.ok_or(ServerInitializeError::ExpectedInitializedNotification(
Some(msg_clone),
))
}
pub async fn serve_server_with_ct<S, T, E, A>(
service: S,
transport: T,
@ -246,12 +230,41 @@ where
ServerInitializeError::transport::<T>(error, "sending initialize response")
})?;
// Wait for initialize notification
let notification = expect_notification(&mut transport, "initialize notification").await?;
let ClientNotification::InitializedNotification(_) = notification else {
return Err(ServerInitializeError::ExpectedInitializedNotification(
Some(ClientJsonRpcMessage::notification(notification)),
));
// Wait for initialized notification. The MCP spec permits logging/setLevel and ping
// before initialized; VS Code sends setLevel immediately after the initialize response.
let notification = loop {
let msg = expect_next_message(&mut transport, "initialize notification").await?;
match msg {
ClientJsonRpcMessage::Notification(n)
if matches!(
n.notification,
ClientNotification::InitializedNotification(_)
) =>
{
break n.notification;
}
ClientJsonRpcMessage::Request(req)
if matches!(
req.request,
ClientRequest::SetLevelRequest(_) | ClientRequest::PingRequest(_)
) =>
{
transport
.send(ServerJsonRpcMessage::response(
ServerResult::EmptyResult(EmptyResult {}),
req.id,
))
.await
.map_err(|error| {
ServerInitializeError::transport::<T>(error, "sending pre-init response")
})?;
}
other => {
return Err(ServerInitializeError::ExpectedInitializedNotification(
Some(other),
));
}
}
};
let context = NotificationContext {
meta: notification.get_meta().clone(),