feat: add support for custom client notifications (#556)
MCP servers, particularly ones that offer "experimental" capabilities, may wish to handle custom client notifications that are not part of the standard MCP specification. This change introduces a new `CustomClientNotification` type that allows a server to process such custom notifications. - introduces `CustomClientNotification` to carry arbitrary methods/params while still preserving meta/extensions; wires it into the `ClientNotification` union and `serde` so `params` can be decoded with `params_as` - allows server handlers to receive custom notifications via a new `on_custom_notification` hook - adds integration coverage that sends a custom client notification end-to-end and asserts the server sees the method and payload Test: ```shell cargo test -p rmcp --features client test_custom_client_notification_reaches_server ```
This commit is contained in:
parent
b6dcb282d9
commit
4ef1a163f4
3 changed files with 263 additions and 7 deletions
|
|
@ -396,6 +396,19 @@
|
|||
"content"
|
||||
]
|
||||
},
|
||||
"CustomClientNotification": {
|
||||
"description": "A catch-all notification the client can use to send custom messages to a server.\n\nThis preserves the raw `method` name and `params` payload so handlers can\ndeserialize them into domain-specific types.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"method": {
|
||||
"type": "string"
|
||||
},
|
||||
"params": true
|
||||
},
|
||||
"required": [
|
||||
"method"
|
||||
]
|
||||
},
|
||||
"ElicitationAction": {
|
||||
"description": "Represents the possible actions a user can take in response to an elicitation request.\n\nWhen a server requests user input through elicitation, the user can:\n- Accept: Provide the requested information and continue\n- Decline: Refuse to provide the information but continue the operation\n- Cancel: Stop the entire operation",
|
||||
"oneOf": [
|
||||
|
|
@ -636,6 +649,9 @@
|
|||
},
|
||||
{
|
||||
"$ref": "#/definitions/NotificationNoParam2"
|
||||
},
|
||||
{
|
||||
"$ref": "#/definitions/CustomClientNotification"
|
||||
}
|
||||
],
|
||||
"required": [
|
||||
|
|
|
|||
|
|
@ -396,6 +396,19 @@
|
|||
"content"
|
||||
]
|
||||
},
|
||||
"CustomClientNotification": {
|
||||
"description": "A catch-all notification the client can use to send custom messages to a server.\n\nThis preserves the raw `method` name and `params` payload so handlers can\ndeserialize them into domain-specific types.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"method": {
|
||||
"type": "string"
|
||||
},
|
||||
"params": true
|
||||
},
|
||||
"required": [
|
||||
"method"
|
||||
]
|
||||
},
|
||||
"ElicitationAction": {
|
||||
"description": "Represents the possible actions a user can take in response to an elicitation request.\n\nWhen a server requests user input through elicitation, the user can:\n- Accept: Provide the requested information and continue\n- Decline: Refuse to provide the information but continue the operation\n- Cancel: Stop the entire operation",
|
||||
"oneOf": [
|
||||
|
|
@ -636,6 +649,9 @@
|
|||
},
|
||||
{
|
||||
"$ref": "#/definitions/NotificationNoParam2"
|
||||
},
|
||||
{
|
||||
"$ref": "#/definitions/CustomClientNotification"
|
||||
}
|
||||
],
|
||||
"required": [
|
||||
|
|
|
|||
|
|
@ -3,10 +3,12 @@ use std::sync::Arc;
|
|||
use rmcp::{
|
||||
ClientHandler, ServerHandler, ServiceExt,
|
||||
model::{
|
||||
ResourceUpdatedNotificationParam, ServerCapabilities, ServerInfo, SubscribeRequestParam,
|
||||
ClientNotification, CustomClientNotification, ResourceUpdatedNotificationParam,
|
||||
ServerCapabilities, ServerInfo, SubscribeRequestParam,
|
||||
},
|
||||
};
|
||||
use tokio::sync::Notify;
|
||||
use serde_json::json;
|
||||
use tokio::sync::{Mutex, Notify};
|
||||
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
|
||||
|
||||
pub struct Server {}
|
||||
|
|
@ -93,3 +95,73 @@ async fn test_server_notification() -> anyhow::Result<()> {
|
|||
client.cancel().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
type CustomNotificationPayload = (String, Option<serde_json::Value>);
|
||||
|
||||
struct CustomServer {
|
||||
receive_signal: Arc<Notify>,
|
||||
payload: Arc<Mutex<Option<CustomNotificationPayload>>>,
|
||||
}
|
||||
|
||||
impl ServerHandler for CustomServer {
|
||||
async fn on_custom_notification(
|
||||
&self,
|
||||
notification: CustomClientNotification,
|
||||
_context: rmcp::service::NotificationContext<rmcp::RoleServer>,
|
||||
) {
|
||||
let CustomClientNotification { method, params, .. } = notification;
|
||||
let mut payload = self.payload.lock().await;
|
||||
*payload = Some((method, params));
|
||||
self.receive_signal.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_custom_client_notification_reaches_server() -> anyhow::Result<()> {
|
||||
let _ = tracing_subscriber::registry()
|
||||
.with(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "debug".to_string().into()),
|
||||
)
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
.try_init();
|
||||
|
||||
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
||||
let receive_signal = Arc::new(Notify::new());
|
||||
let payload = Arc::new(Mutex::new(None));
|
||||
|
||||
{
|
||||
let receive_signal = receive_signal.clone();
|
||||
let payload = payload.clone();
|
||||
tokio::spawn(async move {
|
||||
let server = CustomServer {
|
||||
receive_signal,
|
||||
payload,
|
||||
}
|
||||
.serve(server_transport)
|
||||
.await?;
|
||||
server.waiting().await?;
|
||||
anyhow::Ok(())
|
||||
});
|
||||
}
|
||||
|
||||
let client = ().serve(client_transport).await?;
|
||||
|
||||
client
|
||||
.send_notification(ClientNotification::CustomClientNotification(
|
||||
CustomClientNotification::new(
|
||||
"notifications/custom-test",
|
||||
Some(json!({ "foo": "bar" })),
|
||||
),
|
||||
))
|
||||
.await?;
|
||||
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), receive_signal.notified()).await?;
|
||||
|
||||
let (method, params) = payload.lock().await.clone().expect("payload set");
|
||||
assert_eq!("notifications/custom-test", method);
|
||||
assert_eq!(Some(json!({ "foo": "bar" })), params);
|
||||
|
||||
client.cancel().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue