Skip to content

Commit 2e3cc4a

Browse files
bolinfestbolinfest
authored
feat: add support for custom server notifications (#580)
#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 ```
1 parent 31d242b commit 2e3cc4a

File tree

10 files changed

+180
-35
lines changed

10 files changed

+180
-35
lines changed

‎crates/rmcp/src/handler/client.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,9 @@ impl<H: ClientHandler> Service<RoleClient> for H {
5656
ServerNotification::PromptListChangedNotification(_notification_no_param) => {
5757
self.on_prompt_list_changed(context).await
5858
}
59+
ServerNotification::CustomNotification(notification) => {
60+
self.on_custom_notification(notification, context).await
61+
}
5962
};
6063
Ok(())
6164
}
@@ -166,6 +169,14 @@ pub trait ClientHandler: Sized + Send + Sync + 'static {
166169
) -> impl Future<Output = ()> + Send + '_ {
167170
std::future::ready(())
168171
}
172+
fn on_custom_notification(
173+
&self,
174+
notification: CustomNotification,
175+
context: NotificationContext<RoleClient>,
176+
) -> impl Future<Output = ()> + Send + '_ {
177+
let _ = (notification, context);
178+
std::future::ready(())
179+
}
169180

170181
fn get_info(&self) -> ClientInfo {
171182
ClientInfo::default()

‎crates/rmcp/src/handler/server.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ impl<H: ServerHandler> Service<RoleServer> for H {
9090
ClientNotification::RootsListChangedNotification(_notification) => {
9191
self.on_roots_list_changed(context).await
9292
}
93-
ClientNotification::CustomClientNotification(notification) => {
93+
ClientNotification::CustomNotification(notification) => {
9494
self.on_custom_notification(notification, context).await
9595
}
9696
};
@@ -230,7 +230,7 @@ pub trait ServerHandler: Sized + Send + Sync + 'static {
230230
}
231231
fn on_custom_notification(
232232
&self,
233-
notification: CustomClientNotification,
233+
notification: CustomNotification,
234234
context: NotificationContext<RoleServer>,
235235
) -> impl Future<Output = ()> + Send + '_ {
236236
let _ = (notification, context);

‎crates/rmcp/src/model.rs‎

Lines changed: 39 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -627,13 +627,13 @@ const_string!(CancelledNotificationMethod = "notifications/cancelled");
627627
pub type CancelledNotification =
628628
Notification<CancelledNotificationMethod, CancelledNotificationParam>;
629629

630-
/// A catch-all notification the client can use to send custom messages to a server.
630+
/// A catch-all notification either side can use to send custom messages to its peer.
631631
///
632632
/// This preserves the raw `method` name and `params` payload so handlers can
633633
/// deserialize them into domain-specific types.
634634
#[derive(Debug, Clone)]
635635
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
636-
pub struct CustomClientNotification {
636+
pub struct CustomNotification {
637637
pub method: String,
638638
pub params: Option<Value>,
639639
/// extensions will carry anything possible in the context, including [`Meta`]
@@ -643,7 +643,7 @@ pub struct CustomClientNotification {
643643
pub extensions: Extensions,
644644
}
645645

646-
impl CustomClientNotification {
646+
impl CustomNotification {
647647
pub fn new(method: impl Into<String>, params: Option<Value>) -> Self {
648648
Self {
649649
method: method.into(),
@@ -1786,7 +1786,7 @@ ts_union!(
17861786
| ProgressNotification
17871787
| InitializedNotification
17881788
| RootsListChangedNotification
1789-
| CustomClientNotification;
1789+
| CustomNotification;
17901790
);
17911791

17921792
ts_union!(
@@ -1817,7 +1817,8 @@ ts_union!(
18171817
| ResourceUpdatedNotification
18181818
| ResourceListChangedNotification
18191819
| ToolListChangedNotification
1820-
| PromptListChangedNotification;
1820+
| PromptListChangedNotification
1821+
| CustomNotification;
18211822
);
18221823

18231824
ts_union!(
@@ -1907,7 +1908,7 @@ mod tests {
19071908
serde_json::from_value(raw.clone()).expect("invalid notification");
19081909
match &message {
19091910
ClientJsonRpcMessage::Notification(JsonRpcNotification {
1910-
notification: ClientNotification::CustomClientNotification(notification),
1911+
notification: ClientNotification::CustomNotification(notification),
19111912
..
19121913
}) => {
19131914
assert_eq!(notification.method, "notifications/custom");
@@ -1927,6 +1928,38 @@ mod tests {
19271928
assert_eq!(json, raw);
19281929
}
19291930

1931+
#[test]
1932+
fn test_custom_server_notification_roundtrip() {
1933+
let raw = json!( {
1934+
"jsonrpc": JsonRpcVersion2_0,
1935+
"method": "notifications/custom-server",
1936+
"params": {"hello": "world"},
1937+
});
1938+
1939+
let message: ServerJsonRpcMessage =
1940+
serde_json::from_value(raw.clone()).expect("invalid notification");
1941+
match &message {
1942+
ServerJsonRpcMessage::Notification(JsonRpcNotification {
1943+
notification: ServerNotification::CustomNotification(notification),
1944+
..
1945+
}) => {
1946+
assert_eq!(notification.method, "notifications/custom-server");
1947+
assert_eq!(
1948+
notification
1949+
.params
1950+
.as_ref()
1951+
.and_then(|p| p.get("hello"))
1952+
.expect("hello present"),
1953+
"world"
1954+
);
1955+
}
1956+
_ => panic!("Expected custom server notification"),
1957+
}
1958+
1959+
let json = serde_json::to_value(message).expect("valid json");
1960+
assert_eq!(json, raw);
1961+
}
1962+
19301963
#[test]
19311964
fn test_request_conversion() {
19321965
let raw = json!( {

‎crates/rmcp/src/model/meta.rs‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@ use serde::{Deserialize, Serialize};
44
use serde_json::Value;
55

66
use super::{
7-
ClientNotification, ClientRequest, CustomClientNotification, Extensions, JsonObject,
8-
JsonRpcMessage, NumberOrString, ProgressToken, ServerNotification, ServerRequest,
7+
ClientNotification, ClientRequest, CustomNotification, Extensions, JsonObject, JsonRpcMessage,
8+
NumberOrString, ProgressToken, ServerNotification, ServerRequest,
99
};
1010

1111
pub trait GetMeta {
@@ -18,7 +18,7 @@ pub trait GetExtensions {
1818
fn extensions_mut(&mut self) -> &mut Extensions;
1919
}
2020

21-
impl GetExtensions for CustomClientNotification {
21+
impl GetExtensions for CustomNotification {
2222
fn extensions(&self) -> &Extensions {
2323
&self.extensions
2424
}
@@ -27,7 +27,7 @@ impl GetExtensions for CustomClientNotification {
2727
}
2828
}
2929

30-
impl GetMeta for CustomClientNotification {
30+
impl GetMeta for CustomNotification {
3131
fn get_meta_mut(&mut self) -> &mut Meta {
3232
self.extensions_mut().get_or_insert_default()
3333
}
@@ -104,7 +104,7 @@ variant_extension! {
104104
ProgressNotification
105105
InitializedNotification
106106
RootsListChangedNotification
107-
CustomClientNotification
107+
CustomNotification
108108
}
109109
}
110110

@@ -117,6 +117,7 @@ variant_extension! {
117117
ResourceListChangedNotification
118118
ToolListChangedNotification
119119
PromptListChangedNotification
120+
CustomNotification
120121
}
121122
}
122123
#[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq)]

‎crates/rmcp/src/model/serde_impl.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use std::borrow::Cow;
33
use serde::{Deserialize, Serialize};
44

55
use super::{
6-
CustomClientNotification, Extensions, Meta, Notification, NotificationNoParam, Request,
6+
CustomNotification, Extensions, Meta, Notification, NotificationNoParam, Request,
77
RequestNoParam, RequestOptionalParam,
88
};
99
#[derive(Serialize, Deserialize)]
@@ -249,7 +249,7 @@ where
249249
}
250250
}
251251

252-
impl Serialize for CustomClientNotification {
252+
impl Serialize for CustomNotification {
253253
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
254254
where
255255
S: serde::Serializer,
@@ -277,7 +277,7 @@ impl Serialize for CustomClientNotification {
277277
}
278278
}
279279

280-
impl<'de> Deserialize<'de> for CustomClientNotification {
280+
impl<'de> Deserialize<'de> for CustomNotification {
281281
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
282282
where
283283
D: serde::Deserializer<'de>,
@@ -294,7 +294,7 @@ impl<'de> Deserialize<'de> for CustomClientNotification {
294294
if let Some(meta) = _meta {
295295
extensions.insert(meta);
296296
}
297-
Ok(CustomClientNotification {
297+
Ok(CustomNotification {
298298
extensions,
299299
method: body.method,
300300
params,

‎crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema.json‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -396,8 +396,8 @@
396396
"content"
397397
]
398398
},
399-
"CustomClientNotification": {
400-
"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.",
399+
"CustomNotification": {
400+
"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.",
401401
"type": "object",
402402
"properties": {
403403
"method": {
@@ -651,7 +651,7 @@
651651
"$ref": "#/definitions/NotificationNoParam2"
652652
},
653653
{
654-
"$ref": "#/definitions/CustomClientNotification"
654+
"$ref": "#/definitions/CustomNotification"
655655
}
656656
],
657657
"required": [

‎crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema_current.json‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -396,8 +396,8 @@
396396
"content"
397397
]
398398
},
399-
"CustomClientNotification": {
400-
"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.",
399+
"CustomNotification": {
400+
"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.",
401401
"type": "object",
402402
"properties": {
403403
"method": {
@@ -651,7 +651,7 @@
651651
"$ref": "#/definitions/NotificationNoParam2"
652652
},
653653
{
654-
"$ref": "#/definitions/CustomClientNotification"
654+
"$ref": "#/definitions/CustomNotification"
655655
}
656656
],
657657
"required": [

‎crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema.json‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -392,6 +392,19 @@
392392
"content"
393393
]
394394
},
395+
"CustomNotification": {
396+
"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.",
397+
"type": "object",
398+
"properties": {
399+
"method": {
400+
"type": "string"
401+
},
402+
"params": true
403+
},
404+
"required": [
405+
"method"
406+
]
407+
},
395408
"CancelledNotificationMethod": {
396409
"type": "string",
397410
"format": "const",
@@ -977,6 +990,9 @@
977990
},
978991
{
979992
"$ref": "#/definitions/NotificationNoParam3"
993+
},
994+
{
995+
"$ref": "#/definitions/CustomNotification"
980996
}
981997
],
982998
"required": [

‎crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema_current.json‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -392,6 +392,19 @@
392392
"content"
393393
]
394394
},
395+
"CustomNotification": {
396+
"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.",
397+
"type": "object",
398+
"properties": {
399+
"method": {
400+
"type": "string"
401+
},
402+
"params": true
403+
},
404+
"required": [
405+
"method"
406+
]
407+
},
395408
"CancelledNotificationMethod": {
396409
"type": "string",
397410
"format": "const",
@@ -977,6 +990,9 @@
977990
},
978991
{
979992
"$ref": "#/definitions/NotificationNoParam3"
993+
},
994+
{
995+
"$ref": "#/definitions/CustomNotification"
980996
}
981997
],
982998
"required": [

0 commit comments

Comments
 (0)