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:
parent
9b507f5018
commit
8e5ebb4f5c
1 changed files with 196 additions and 23 deletions
|
|
@ -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(),
|
||||
|
|
|
|||
Loading…
Reference in a new issue