modelcontextprotocol/rust-sdk
Watch
1
0
Fork
You've already forked rust-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

feat: add support for custom server notifications (#580)

Browse source 
https://github.com/modelcontextprotocol/rust-sdk/pull/556 introduced support for
custom client notifications, so this PR makes the complementary change, adding
support for custom server notifications.

MCP clients, particularly ones that offer "experimental" capabilities,
may wish to handle custom server notifications that are not part of the
standard MCP specification. This change introduces a new
`CustomServerNotification` type that allows a client to process
such custom notifications.

- introduces `CustomServerNotification` to carry arbitrary methods/params while
  still preserving meta/extensions; wires it into the `ServerNotification` union
  and `serde` so `params` can be decoded with `params_as`
- allows client handlers to receive custom notifications via a new
  `on_custom_notification` hook
- adds integration coverage that sends a custom server notification end-to-end
  and asserts the client sees the method and payload

Test:

```shell
cargo test -p rmcp --features client test_custom_server_notification_reaches_client
```
This commit is contained in:
Michael Bolin 2025-12-16 09:48:03 -08:00 • committed by GitHub
commit 2e3cc4a973
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 180 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

6
crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema.json
View file

@ -396,8 +396,8 @@
"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.",
"CustomNotification": {
"description": "A catch-all notification either side can use to send custom messages to its peer.\n\nThis preserves the raw `method` name and `params` payload so handlers can\ndeserialize them into domain-specific types.",
"type": "object",
"properties": {
"method": {
@ -651,7 +651,7 @@
"$ref": "#/definitions/NotificationNoParam2"
},
{
"$ref": "#/definitions/CustomClientNotification"
"$ref": "#/definitions/CustomNotification"
}
],
"required": [

6
crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema_current.json
View file

@ -396,8 +396,8 @@
"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.",
"CustomNotification": {
"description": "A catch-all notification either side can use to send custom messages to its peer.\n\nThis preserves the raw `method` name and `params` payload so handlers can\ndeserialize them into domain-specific types.",
"type": "object",
"properties": {
"method": {
@ -651,7 +651,7 @@
"$ref": "#/definitions/NotificationNoParam2"
},
{
"$ref": "#/definitions/CustomClientNotification"
"$ref": "#/definitions/CustomNotification"
}
],
"required": [

16
crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema.json
View file

@ -392,6 +392,19 @@
"content"
]
},
"CustomNotification": {
"description": "A catch-all notification either side can use to send custom messages to its peer.\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"
]
},
"CancelledNotificationMethod": {
"type": "string",
"format": "const",
@ -977,6 +990,9 @@
},
{
"$ref": "#/definitions/NotificationNoParam3"
},
{
"$ref": "#/definitions/CustomNotification"
}
],
"required": [

16
crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema_current.json
View file

@ -392,6 +392,19 @@
"content"
]
},
"CustomNotification": {
"description": "A catch-all notification either side can use to send custom messages to its peer.\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"
]
},
"CancelledNotificationMethod": {
"type": "string",
"format": "const",
@ -977,6 +990,9 @@
},
{
"$ref": "#/definitions/NotificationNoParam3"
},
{
"$ref": "#/definitions/CustomNotification"
}
],
"required": [

92
crates/rmcp/tests/test_notification.rs
View file

@ -3,8 +3,8 @@ use std::sync::Arc;
use rmcp::{
ClientHandler, ServerHandler, ServiceExt,
model::{
ClientNotification, CustomClientNotification, ResourceUpdatedNotificationParam,
ServerCapabilities, ServerInfo, SubscribeRequestParam,
ClientNotification, CustomNotification, ResourceUpdatedNotificationParam,
ServerCapabilities, ServerInfo, ServerNotification, SubscribeRequestParam,
},
};
use serde_json::json;
@ -106,12 +106,11 @@ struct CustomServer {
impl ServerHandler for CustomServer {
async fn on_custom_notification(
&self,
notification: CustomClientNotification,
notification: CustomNotification,
_context: rmcp::service::NotificationContext<rmcp::RoleServer>,
) {
let CustomClientNotification { method, params, .. } = notification;
let mut payload = self.payload.lock().await;
*payload = Some((method, params));
let CustomNotification { method, params, .. } = notification;
*self.payload.lock().await = Some((method, params));
self.receive_signal.notify_one();
}
}
@ -148,20 +147,89 @@ async fn test_custom_client_notification_reaches_server() -> anyhow::Result<()>
let client = ().serve(client_transport).await?;
client
.send_notification(ClientNotification::CustomClientNotification(
CustomClientNotification::new(
"notifications/custom-test",
Some(json!({ "foo": "bar" })),
),
.send_notification(ClientNotification::CustomNotification(
CustomNotification::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");
let (method, params) = payload.lock().await.take().expect("payload set");
assert_eq!("notifications/custom-test", method);
assert_eq!(Some(json!({ "foo": "bar" })), params);
client.cancel().await?;
Ok(())
}
struct CustomServerNotifier;
impl ServerHandler for CustomServerNotifier {
async fn on_initialized(&self, context: rmcp::service::NotificationContext<rmcp::RoleServer>) {
let peer = context.peer.clone();
tokio::spawn(async move {
peer.send_notification(ServerNotification::CustomNotification(
CustomNotification::new(
"notifications/custom-test",
Some(json!({ "hello": "world" })),
),
))
.await
.expect("send custom notification");
});
}
}
struct CustomClient {
receive_signal: Arc<Notify>,
payload: Arc<Mutex<Option<CustomNotificationPayload>>>,
}
impl ClientHandler for CustomClient {
async fn on_custom_notification(
&self,
notification: CustomNotification,
_context: rmcp::service::NotificationContext<rmcp::RoleClient>,
) {
let CustomNotification { method, params, .. } = notification;
*self.payload.lock().await = Some((method, params));
self.receive_signal.notify_one();
}
}
#[tokio::test]
async fn test_custom_server_notification_reaches_client() -> 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);
tokio::spawn(async move {
let server = CustomServerNotifier {}.serve(server_transport).await?;
server.waiting().await?;
anyhow::Ok(())
});
let receive_signal = Arc::new(Notify::new());
let payload = Arc::new(Mutex::new(None));
let client = CustomClient {
receive_signal: receive_signal.clone(),
payload: payload.clone(),
}
.serve(client_transport)
.await?;
tokio::time::timeout(std::time::Duration::from_secs(5), receive_signal.notified()).await?;
let (method, params) = payload.lock().await.take().expect("payload set");
assert_eq!("notifications/custom-test", method);
assert_eq!(Some(json!({ "hello": "world" })), params);
client.cancel().await?;
Ok(())
}