feat: add support for custom requests (#590)
#580 and #556 introduced support for custom notifications, so this PR takes the next logical step and adds support for custom requests: - Introduces `CustomRequest` and `CustomResult` model types, wires them into the client/server request and result unions, and allows `ClientRequest::method()` to return the dynamic method name. - Implements serde and meta handling for `CustomRequest` so `_meta` is carried through extensions; adds default `on_custom_request` handlers that return `METHOD_NOT_FOUND` unless overridden. - Updates JSON schema fixtures to include the new request/result shapes and `EmptyObject` strictness. - Adds tests for custom request roundtrips and end-to-end client↔server handling. - Focused integration test in `crates/rmcp/tests/test_custom_request.rs`. For additional testing, I used this locally to update Codex to use a custom request instead of a custom notification so that it gets an "ack" from the MCP server to ensure it has processed the update before sending more messages: https://github.com/openai/codex/pull/8142.
This commit is contained in:
parent
2e3cc4a973
commit
e0faf1ed51
10 changed files with 524 additions and 38 deletions
|
|
@ -26,6 +26,10 @@ impl<H: ClientHandler> Service<RoleClient> for H {
|
||||||
.create_elicitation(request.params, context)
|
.create_elicitation(request.params, context)
|
||||||
.await
|
.await
|
||||||
.map(ClientResult::CreateElicitationResult),
|
.map(ClientResult::CreateElicitationResult),
|
||||||
|
ServerRequest::CustomRequest(request) => self
|
||||||
|
.on_custom_request(request, context)
|
||||||
|
.await
|
||||||
|
.map(ClientResult::CustomResult),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -123,6 +127,20 @@ pub trait ClientHandler: Sized + Send + Sync + 'static {
|
||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn on_custom_request(
|
||||||
|
&self,
|
||||||
|
request: CustomRequest,
|
||||||
|
context: RequestContext<RoleClient>,
|
||||||
|
) -> impl Future<Output = Result<CustomResult, McpError>> + Send + '_ {
|
||||||
|
let CustomRequest { method, .. } = request;
|
||||||
|
let _ = context;
|
||||||
|
std::future::ready(Err(McpError::new(
|
||||||
|
ErrorCode::METHOD_NOT_FOUND,
|
||||||
|
method,
|
||||||
|
None,
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
|
||||||
fn on_cancelled(
|
fn on_cancelled(
|
||||||
&self,
|
&self,
|
||||||
params: CancelledNotificationParam,
|
params: CancelledNotificationParam,
|
||||||
|
|
|
||||||
|
|
@ -69,6 +69,10 @@ impl<H: ServerHandler> Service<RoleServer> for H {
|
||||||
.list_tools(request.params, context)
|
.list_tools(request.params, context)
|
||||||
.await
|
.await
|
||||||
.map(ServerResult::ListToolsResult),
|
.map(ServerResult::ListToolsResult),
|
||||||
|
ClientRequest::CustomRequest(request) => self
|
||||||
|
.on_custom_request(request, context)
|
||||||
|
.await
|
||||||
|
.map(ServerResult::CustomResult),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -200,6 +204,19 @@ pub trait ServerHandler: Sized + Send + Sync + 'static {
|
||||||
) -> impl Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
|
) -> impl Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
|
||||||
std::future::ready(Ok(ListToolsResult::default()))
|
std::future::ready(Ok(ListToolsResult::default()))
|
||||||
}
|
}
|
||||||
|
fn on_custom_request(
|
||||||
|
&self,
|
||||||
|
request: CustomRequest,
|
||||||
|
context: RequestContext<RoleServer>,
|
||||||
|
) -> impl Future<Output = Result<CustomResult, McpError>> + Send + '_ {
|
||||||
|
let CustomRequest { method, .. } = request;
|
||||||
|
let _ = context;
|
||||||
|
std::future::ready(Err(McpError::new(
|
||||||
|
ErrorCode::METHOD_NOT_FOUND,
|
||||||
|
method,
|
||||||
|
None,
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
|
||||||
fn on_cancelled(
|
fn on_cancelled(
|
||||||
&self,
|
&self,
|
||||||
|
|
|
||||||
|
|
@ -55,6 +55,7 @@ macro_rules! object {
|
||||||
///
|
///
|
||||||
/// without returning any specific data.
|
/// without returning any specific data.
|
||||||
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Copy, Eq)]
|
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Copy, Eq)]
|
||||||
|
#[serde(deny_unknown_fields)]
|
||||||
#[cfg_attr(feature = "server", derive(schemars::JsonSchema))]
|
#[cfg_attr(feature = "server", derive(schemars::JsonSchema))]
|
||||||
pub struct EmptyObject {}
|
pub struct EmptyObject {}
|
||||||
|
|
||||||
|
|
@ -606,6 +607,23 @@ impl From<EmptyResult> for () {
|
||||||
fn from(_value: EmptyResult) {}
|
fn from(_value: EmptyResult) {}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// A catch-all response either side can use for custom requests.
|
||||||
|
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
|
||||||
|
#[serde(transparent)]
|
||||||
|
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
|
||||||
|
pub struct CustomResult(pub Value);
|
||||||
|
|
||||||
|
impl CustomResult {
|
||||||
|
pub fn new(result: Value) -> Self {
|
||||||
|
Self(result)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Deserialize the result into a strongly-typed structure.
|
||||||
|
pub fn result_as<T: DeserializeOwned>(&self) -> Result<T, serde_json::Error> {
|
||||||
|
serde_json::from_value(self.0.clone())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
|
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
|
||||||
#[serde(rename_all = "camelCase")]
|
#[serde(rename_all = "camelCase")]
|
||||||
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
|
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
|
||||||
|
|
@ -661,6 +679,40 @@ impl CustomNotification {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// A catch-all request either side can use to send custom messages to its peer.
|
||||||
|
///
|
||||||
|
/// This preserves the raw `method` name and `params` payload so handlers can
|
||||||
|
/// deserialize them into domain-specific types.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
|
||||||
|
pub struct CustomRequest {
|
||||||
|
pub method: String,
|
||||||
|
pub params: Option<Value>,
|
||||||
|
/// extensions will carry anything possible in the context, including [`Meta`]
|
||||||
|
///
|
||||||
|
/// this is similar with the Extensions in `http` crate
|
||||||
|
#[cfg_attr(feature = "schemars", schemars(skip))]
|
||||||
|
pub extensions: Extensions,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CustomRequest {
|
||||||
|
pub fn new(method: impl Into<String>, params: Option<Value>) -> Self {
|
||||||
|
Self {
|
||||||
|
method: method.into(),
|
||||||
|
params,
|
||||||
|
extensions: Extensions::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Deserialize `params` into a strongly-typed structure.
|
||||||
|
pub fn params_as<T: DeserializeOwned>(&self) -> Result<Option<T>, serde_json::Error> {
|
||||||
|
self.params
|
||||||
|
.as_ref()
|
||||||
|
.map(|params| serde_json::from_value(params.clone()))
|
||||||
|
.transpose()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
const_string!(InitializeResultMethod = "initialize");
|
const_string!(InitializeResultMethod = "initialize");
|
||||||
/// # Initialization
|
/// # Initialization
|
||||||
/// This request is sent from the client to the server when it first connects, asking it to begin initialization.
|
/// This request is sent from the client to the server when it first connects, asking it to begin initialization.
|
||||||
|
|
@ -1757,11 +1809,12 @@ ts_union!(
|
||||||
| SubscribeRequest
|
| SubscribeRequest
|
||||||
| UnsubscribeRequest
|
| UnsubscribeRequest
|
||||||
| CallToolRequest
|
| CallToolRequest
|
||||||
| ListToolsRequest;
|
| ListToolsRequest
|
||||||
|
| CustomRequest;
|
||||||
);
|
);
|
||||||
|
|
||||||
impl ClientRequest {
|
impl ClientRequest {
|
||||||
pub fn method(&self) -> &'static str {
|
pub fn method(&self) -> &str {
|
||||||
match &self {
|
match &self {
|
||||||
ClientRequest::PingRequest(r) => r.method.as_str(),
|
ClientRequest::PingRequest(r) => r.method.as_str(),
|
||||||
ClientRequest::InitializeRequest(r) => r.method.as_str(),
|
ClientRequest::InitializeRequest(r) => r.method.as_str(),
|
||||||
|
|
@ -1776,6 +1829,7 @@ impl ClientRequest {
|
||||||
ClientRequest::UnsubscribeRequest(r) => r.method.as_str(),
|
ClientRequest::UnsubscribeRequest(r) => r.method.as_str(),
|
||||||
ClientRequest::CallToolRequest(r) => r.method.as_str(),
|
ClientRequest::CallToolRequest(r) => r.method.as_str(),
|
||||||
ClientRequest::ListToolsRequest(r) => r.method.as_str(),
|
ClientRequest::ListToolsRequest(r) => r.method.as_str(),
|
||||||
|
ClientRequest::CustomRequest(r) => r.method.as_str(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -1790,7 +1844,12 @@ ts_union!(
|
||||||
);
|
);
|
||||||
|
|
||||||
ts_union!(
|
ts_union!(
|
||||||
export type ClientResult = box CreateMessageResult | ListRootsResult | CreateElicitationResult | EmptyResult;
|
export type ClientResult =
|
||||||
|
box CreateMessageResult
|
||||||
|
| ListRootsResult
|
||||||
|
| CreateElicitationResult
|
||||||
|
| EmptyResult
|
||||||
|
| CustomResult;
|
||||||
);
|
);
|
||||||
|
|
||||||
impl ClientResult {
|
impl ClientResult {
|
||||||
|
|
@ -1806,7 +1865,8 @@ ts_union!(
|
||||||
| PingRequest
|
| PingRequest
|
||||||
| CreateMessageRequest
|
| CreateMessageRequest
|
||||||
| ListRootsRequest
|
| ListRootsRequest
|
||||||
| CreateElicitationRequest;
|
| CreateElicitationRequest
|
||||||
|
| CustomRequest;
|
||||||
);
|
);
|
||||||
|
|
||||||
ts_union!(
|
ts_union!(
|
||||||
|
|
@ -1834,6 +1894,7 @@ ts_union!(
|
||||||
| ListToolsResult
|
| ListToolsResult
|
||||||
| CreateElicitationResult
|
| CreateElicitationResult
|
||||||
| EmptyResult
|
| EmptyResult
|
||||||
|
| CustomResult
|
||||||
;
|
;
|
||||||
);
|
);
|
||||||
|
|
||||||
|
|
@ -1960,6 +2021,40 @@ mod tests {
|
||||||
assert_eq!(json, raw);
|
assert_eq!(json, raw);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_custom_request_roundtrip() {
|
||||||
|
let raw = json!( {
|
||||||
|
"jsonrpc": JsonRpcVersion2_0,
|
||||||
|
"id": 42,
|
||||||
|
"method": "requests/custom",
|
||||||
|
"params": {"foo": "bar"},
|
||||||
|
});
|
||||||
|
|
||||||
|
let message: ClientJsonRpcMessage =
|
||||||
|
serde_json::from_value(raw.clone()).expect("invalid request");
|
||||||
|
match &message {
|
||||||
|
ClientJsonRpcMessage::Request(JsonRpcRequest { id, request, .. }) => {
|
||||||
|
assert_eq!(id, &RequestId::Number(42));
|
||||||
|
match request {
|
||||||
|
ClientRequest::CustomRequest(custom) => {
|
||||||
|
let expected_request = json!({
|
||||||
|
"method": "requests/custom",
|
||||||
|
"params": {"foo": "bar"},
|
||||||
|
});
|
||||||
|
let actual_request =
|
||||||
|
serde_json::to_value(custom).expect("serialize custom request");
|
||||||
|
assert_eq!(actual_request, expected_request);
|
||||||
|
}
|
||||||
|
other => panic!("Expected custom request, got: {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
other => panic!("Expected request, got: {other:?}"),
|
||||||
|
}
|
||||||
|
|
||||||
|
let json = serde_json::to_value(message).expect("valid json");
|
||||||
|
assert_eq!(json, raw);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_request_conversion() {
|
fn test_request_conversion() {
|
||||||
let raw = json!( {
|
let raw = json!( {
|
||||||
|
|
|
||||||
|
|
@ -4,8 +4,8 @@ use serde::{Deserialize, Serialize};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
ClientNotification, ClientRequest, CustomNotification, Extensions, JsonObject, JsonRpcMessage,
|
ClientNotification, ClientRequest, CustomNotification, CustomRequest, Extensions, JsonObject,
|
||||||
NumberOrString, ProgressToken, ServerNotification, ServerRequest,
|
JsonRpcMessage, NumberOrString, ProgressToken, ServerNotification, ServerRequest,
|
||||||
};
|
};
|
||||||
|
|
||||||
pub trait GetMeta {
|
pub trait GetMeta {
|
||||||
|
|
@ -38,6 +38,26 @@ impl GetMeta for CustomNotification {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl GetExtensions for CustomRequest {
|
||||||
|
fn extensions(&self) -> &Extensions {
|
||||||
|
&self.extensions
|
||||||
|
}
|
||||||
|
fn extensions_mut(&mut self) -> &mut Extensions {
|
||||||
|
&mut self.extensions
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl GetMeta for CustomRequest {
|
||||||
|
fn get_meta_mut(&mut self) -> &mut Meta {
|
||||||
|
self.extensions_mut().get_or_insert_default()
|
||||||
|
}
|
||||||
|
fn get_meta(&self) -> &Meta {
|
||||||
|
self.extensions()
|
||||||
|
.get::<Meta>()
|
||||||
|
.unwrap_or(Meta::static_empty())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
macro_rules! variant_extension {
|
macro_rules! variant_extension {
|
||||||
(
|
(
|
||||||
$Enum: ident {
|
$Enum: ident {
|
||||||
|
|
@ -86,6 +106,7 @@ variant_extension! {
|
||||||
UnsubscribeRequest
|
UnsubscribeRequest
|
||||||
CallToolRequest
|
CallToolRequest
|
||||||
ListToolsRequest
|
ListToolsRequest
|
||||||
|
CustomRequest
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -95,6 +116,7 @@ variant_extension! {
|
||||||
CreateMessageRequest
|
CreateMessageRequest
|
||||||
ListRootsRequest
|
ListRootsRequest
|
||||||
CreateElicitationRequest
|
CreateElicitationRequest
|
||||||
|
CustomRequest
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -3,8 +3,8 @@ use std::borrow::Cow;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
CustomNotification, Extensions, Meta, Notification, NotificationNoParam, Request,
|
CustomNotification, CustomRequest, Extensions, Meta, Notification, NotificationNoParam,
|
||||||
RequestNoParam, RequestOptionalParam,
|
Request, RequestNoParam, RequestOptionalParam,
|
||||||
};
|
};
|
||||||
#[derive(Serialize, Deserialize)]
|
#[derive(Serialize, Deserialize)]
|
||||||
struct WithMeta<'a, P> {
|
struct WithMeta<'a, P> {
|
||||||
|
|
@ -249,6 +249,59 @@ where
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl Serialize for CustomRequest {
|
||||||
|
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
||||||
|
where
|
||||||
|
S: serde::Serializer,
|
||||||
|
{
|
||||||
|
let extensions = &self.extensions;
|
||||||
|
let _meta = extensions.get::<Meta>().map(Cow::Borrowed);
|
||||||
|
let params = self.params.as_ref();
|
||||||
|
|
||||||
|
let params = if _meta.is_some() || params.is_some() {
|
||||||
|
Some(WithMeta {
|
||||||
|
_meta,
|
||||||
|
_rest: &self.params,
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
|
ProxyOptionalParam::serialize(
|
||||||
|
&ProxyOptionalParam {
|
||||||
|
method: &self.method,
|
||||||
|
params,
|
||||||
|
},
|
||||||
|
serializer,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'de> Deserialize<'de> for CustomRequest {
|
||||||
|
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
||||||
|
where
|
||||||
|
D: serde::Deserializer<'de>,
|
||||||
|
{
|
||||||
|
let body =
|
||||||
|
ProxyOptionalParam::<'_, _, Option<serde_json::Value>>::deserialize(deserializer)?;
|
||||||
|
let mut params = None;
|
||||||
|
let mut _meta = None;
|
||||||
|
if let Some(body_params) = body.params {
|
||||||
|
params = body_params._rest;
|
||||||
|
_meta = body_params._meta.map(|m| m.into_owned());
|
||||||
|
}
|
||||||
|
let mut extensions = Extensions::new();
|
||||||
|
if let Some(meta) = _meta {
|
||||||
|
extensions.insert(meta);
|
||||||
|
}
|
||||||
|
Ok(CustomRequest {
|
||||||
|
extensions,
|
||||||
|
method: body.method,
|
||||||
|
params,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
impl Serialize for CustomNotification {
|
impl Serialize for CustomNotification {
|
||||||
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
||||||
where
|
where
|
||||||
|
|
|
||||||
189
crates/rmcp/tests/test_custom_request.rs
Normal file
189
crates/rmcp/tests/test_custom_request.rs
Normal file
|
|
@ -0,0 +1,189 @@
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use rmcp::{
|
||||||
|
ClientHandler, ServerHandler, ServiceExt,
|
||||||
|
model::{
|
||||||
|
ClientRequest, ClientResult, CustomRequest, CustomResult, ServerRequest, ServerResult,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
use serde_json::json;
|
||||||
|
use tokio::sync::{Mutex, Notify};
|
||||||
|
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
|
||||||
|
|
||||||
|
type CustomRequestPayload = (String, Option<serde_json::Value>);
|
||||||
|
|
||||||
|
struct CustomRequestServer {
|
||||||
|
receive_signal: Arc<Notify>,
|
||||||
|
payload: Arc<Mutex<Option<CustomRequestPayload>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ServerHandler for CustomRequestServer {
|
||||||
|
async fn on_custom_request(
|
||||||
|
&self,
|
||||||
|
request: CustomRequest,
|
||||||
|
_context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
||||||
|
) -> Result<CustomResult, rmcp::ErrorData> {
|
||||||
|
let CustomRequest { method, params, .. } = request;
|
||||||
|
*self.payload.lock().await = Some((method, params));
|
||||||
|
self.receive_signal.notify_one();
|
||||||
|
Ok(CustomResult::new(json!({ "status": "ok" })))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_custom_client_request_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 = CustomRequestServer {
|
||||||
|
receive_signal,
|
||||||
|
payload,
|
||||||
|
}
|
||||||
|
.serve(server_transport)
|
||||||
|
.await?;
|
||||||
|
server.waiting().await?;
|
||||||
|
anyhow::Ok(())
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let client = ().serve(client_transport).await?;
|
||||||
|
|
||||||
|
let response = client
|
||||||
|
.send_request(ClientRequest::CustomRequest(CustomRequest::new(
|
||||||
|
"requests/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.take().expect("payload set");
|
||||||
|
assert_eq!("requests/custom-test", method);
|
||||||
|
assert_eq!(Some(json!({ "foo": "bar" })), params);
|
||||||
|
|
||||||
|
match response {
|
||||||
|
ServerResult::CustomResult(result) => {
|
||||||
|
assert_eq!(result.0, json!({ "status": "ok" }));
|
||||||
|
}
|
||||||
|
other => panic!("Expected custom result, got: {other:?}"),
|
||||||
|
}
|
||||||
|
|
||||||
|
client.cancel().await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
struct CustomRequestClient {
|
||||||
|
receive_signal: Arc<Notify>,
|
||||||
|
payload: Arc<Mutex<Option<CustomRequestPayload>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ClientHandler for CustomRequestClient {
|
||||||
|
async fn on_custom_request(
|
||||||
|
&self,
|
||||||
|
request: CustomRequest,
|
||||||
|
_context: rmcp::service::RequestContext<rmcp::RoleClient>,
|
||||||
|
) -> Result<CustomResult, rmcp::ErrorData> {
|
||||||
|
let CustomRequest { method, params, .. } = request;
|
||||||
|
*self.payload.lock().await = Some((method, params));
|
||||||
|
self.receive_signal.notify_one();
|
||||||
|
Ok(CustomResult::new(json!({ "status": "ok" })))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct CustomRequestServerNotifier {
|
||||||
|
receive_signal: Arc<Notify>,
|
||||||
|
response: Arc<Mutex<Option<Result<serde_json::Value, String>>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ServerHandler for CustomRequestServerNotifier {
|
||||||
|
async fn on_initialized(&self, context: rmcp::service::NotificationContext<rmcp::RoleServer>) {
|
||||||
|
let peer = context.peer.clone();
|
||||||
|
let receive_signal = self.receive_signal.clone();
|
||||||
|
let response = self.response.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let result = peer
|
||||||
|
.send_request(ServerRequest::CustomRequest(CustomRequest::new(
|
||||||
|
"requests/custom-server",
|
||||||
|
Some(json!({ "ping": "pong" })),
|
||||||
|
)))
|
||||||
|
.await;
|
||||||
|
let payload = match result {
|
||||||
|
Ok(ClientResult::CustomResult(result)) => Ok(result.0),
|
||||||
|
Ok(other) => Err(format!("Unexpected response: {other:?}")),
|
||||||
|
Err(err) => Err(format!("Failed to send request: {err:?}")),
|
||||||
|
};
|
||||||
|
*response.lock().await = Some(payload);
|
||||||
|
receive_signal.notify_one();
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_custom_server_request_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);
|
||||||
|
let response_signal = Arc::new(Notify::new());
|
||||||
|
let response = Arc::new(Mutex::new(None));
|
||||||
|
tokio::spawn({
|
||||||
|
let response_signal = response_signal.clone();
|
||||||
|
let response = response.clone();
|
||||||
|
async move {
|
||||||
|
let server = CustomRequestServerNotifier {
|
||||||
|
receive_signal: response_signal,
|
||||||
|
response,
|
||||||
|
}
|
||||||
|
.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 = CustomRequestClient {
|
||||||
|
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?;
|
||||||
|
tokio::time::timeout(
|
||||||
|
std::time::Duration::from_secs(5),
|
||||||
|
response_signal.notified(),
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let (method, params) = payload.lock().await.take().expect("payload set");
|
||||||
|
assert_eq!("requests/custom-server", method);
|
||||||
|
assert_eq!(Some(json!({ "ping": "pong" })), params);
|
||||||
|
|
||||||
|
let response = response.lock().await.take().expect("response set");
|
||||||
|
let response = response.expect("custom request response ok");
|
||||||
|
assert_eq!(response, json!({ "status": "ok" }));
|
||||||
|
|
||||||
|
client.cancel().await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
@ -288,6 +288,9 @@
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"$ref": "#/definitions/EmptyObject"
|
"$ref": "#/definitions/EmptyObject"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"$ref": "#/definitions/CustomResult"
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
|
@ -409,6 +412,22 @@
|
||||||
"method"
|
"method"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
"CustomRequest": {
|
||||||
|
"description": "A catch-all request 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"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"CustomResult": {
|
||||||
|
"description": "A catch-all response either side can use for custom requests."
|
||||||
|
},
|
||||||
"ElicitationAction": {
|
"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",
|
"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": [
|
"oneOf": [
|
||||||
|
|
@ -444,7 +463,8 @@
|
||||||
},
|
},
|
||||||
"EmptyObject": {
|
"EmptyObject": {
|
||||||
"description": "This is commonly used for representing empty objects in MCP messages.\n\nwithout returning any specific data.",
|
"description": "This is commonly used for representing empty objects in MCP messages.\n\nwithout returning any specific data.",
|
||||||
"type": "object"
|
"type": "object",
|
||||||
|
"additionalProperties": false
|
||||||
},
|
},
|
||||||
"ErrorCode": {
|
"ErrorCode": {
|
||||||
"description": "Standard JSON-RPC error codes used throughout the MCP protocol.\n\nThese codes follow the JSON-RPC 2.0 specification and provide\nstandardized error reporting across all MCP implementations.",
|
"description": "Standard JSON-RPC error codes used throughout the MCP protocol.\n\nThese codes follow the JSON-RPC 2.0 specification and provide\nstandardized error reporting across all MCP implementations.",
|
||||||
|
|
@ -707,6 +727,9 @@
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"$ref": "#/definitions/RequestOptionalParam4"
|
"$ref": "#/definitions/RequestOptionalParam4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"$ref": "#/definitions/CustomRequest"
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
"required": [
|
"required": [
|
||||||
|
|
|
||||||
|
|
@ -288,6 +288,9 @@
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"$ref": "#/definitions/EmptyObject"
|
"$ref": "#/definitions/EmptyObject"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"$ref": "#/definitions/CustomResult"
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
|
@ -409,6 +412,22 @@
|
||||||
"method"
|
"method"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
"CustomRequest": {
|
||||||
|
"description": "A catch-all request 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"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"CustomResult": {
|
||||||
|
"description": "A catch-all response either side can use for custom requests."
|
||||||
|
},
|
||||||
"ElicitationAction": {
|
"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",
|
"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": [
|
"oneOf": [
|
||||||
|
|
@ -444,7 +463,8 @@
|
||||||
},
|
},
|
||||||
"EmptyObject": {
|
"EmptyObject": {
|
||||||
"description": "This is commonly used for representing empty objects in MCP messages.\n\nwithout returning any specific data.",
|
"description": "This is commonly used for representing empty objects in MCP messages.\n\nwithout returning any specific data.",
|
||||||
"type": "object"
|
"type": "object",
|
||||||
|
"additionalProperties": false
|
||||||
},
|
},
|
||||||
"ErrorCode": {
|
"ErrorCode": {
|
||||||
"description": "Standard JSON-RPC error codes used throughout the MCP protocol.\n\nThese codes follow the JSON-RPC 2.0 specification and provide\nstandardized error reporting across all MCP implementations.",
|
"description": "Standard JSON-RPC error codes used throughout the MCP protocol.\n\nThese codes follow the JSON-RPC 2.0 specification and provide\nstandardized error reporting across all MCP implementations.",
|
||||||
|
|
@ -707,6 +727,9 @@
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"$ref": "#/definitions/RequestOptionalParam4"
|
"$ref": "#/definitions/RequestOptionalParam4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"$ref": "#/definitions/CustomRequest"
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
"required": [
|
"required": [
|
||||||
|
|
|
||||||
|
|
@ -392,19 +392,6 @@
|
||||||
"content"
|
"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": {
|
"CancelledNotificationMethod": {
|
||||||
"type": "string",
|
"type": "string",
|
||||||
"format": "const",
|
"format": "const",
|
||||||
|
|
@ -606,6 +593,35 @@
|
||||||
"maxTokens"
|
"maxTokens"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
"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"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"CustomRequest": {
|
||||||
|
"description": "A catch-all request 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"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"CustomResult": {
|
||||||
|
"description": "A catch-all response either side can use for custom requests."
|
||||||
|
},
|
||||||
"ElicitationAction": {
|
"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",
|
"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": [
|
"oneOf": [
|
||||||
|
|
@ -682,7 +698,8 @@
|
||||||
},
|
},
|
||||||
"EmptyObject": {
|
"EmptyObject": {
|
||||||
"description": "This is commonly used for representing empty objects in MCP messages.\n\nwithout returning any specific data.",
|
"description": "This is commonly used for representing empty objects in MCP messages.\n\nwithout returning any specific data.",
|
||||||
"type": "object"
|
"type": "object",
|
||||||
|
"additionalProperties": false
|
||||||
},
|
},
|
||||||
"EnumSchema": {
|
"EnumSchema": {
|
||||||
"description": "Schema definition for enum properties.\n\nCompliant with MCP 2025-06-18 specification for elicitation schemas.\nEnums must have string type and can optionally include human-readable names.",
|
"description": "Schema definition for enum properties.\n\nCompliant with MCP 2025-06-18 specification for elicitation schemas.\nEnums must have string type and can optionally include human-readable names.",
|
||||||
|
|
@ -1021,6 +1038,9 @@
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"$ref": "#/definitions/Request2"
|
"$ref": "#/definitions/Request2"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"$ref": "#/definitions/CustomRequest"
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
"required": [
|
"required": [
|
||||||
|
|
@ -2271,6 +2291,9 @@
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"$ref": "#/definitions/EmptyObject"
|
"$ref": "#/definitions/EmptyObject"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"$ref": "#/definitions/CustomResult"
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
|
|
||||||
|
|
@ -392,19 +392,6 @@
|
||||||
"content"
|
"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": {
|
"CancelledNotificationMethod": {
|
||||||
"type": "string",
|
"type": "string",
|
||||||
"format": "const",
|
"format": "const",
|
||||||
|
|
@ -606,6 +593,35 @@
|
||||||
"maxTokens"
|
"maxTokens"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
"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"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"CustomRequest": {
|
||||||
|
"description": "A catch-all request 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"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"CustomResult": {
|
||||||
|
"description": "A catch-all response either side can use for custom requests."
|
||||||
|
},
|
||||||
"ElicitationAction": {
|
"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",
|
"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": [
|
"oneOf": [
|
||||||
|
|
@ -682,7 +698,8 @@
|
||||||
},
|
},
|
||||||
"EmptyObject": {
|
"EmptyObject": {
|
||||||
"description": "This is commonly used for representing empty objects in MCP messages.\n\nwithout returning any specific data.",
|
"description": "This is commonly used for representing empty objects in MCP messages.\n\nwithout returning any specific data.",
|
||||||
"type": "object"
|
"type": "object",
|
||||||
|
"additionalProperties": false
|
||||||
},
|
},
|
||||||
"EnumSchema": {
|
"EnumSchema": {
|
||||||
"description": "Schema definition for enum properties.\n\nCompliant with MCP 2025-06-18 specification for elicitation schemas.\nEnums must have string type and can optionally include human-readable names.",
|
"description": "Schema definition for enum properties.\n\nCompliant with MCP 2025-06-18 specification for elicitation schemas.\nEnums must have string type and can optionally include human-readable names.",
|
||||||
|
|
@ -1021,6 +1038,9 @@
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"$ref": "#/definitions/Request2"
|
"$ref": "#/definitions/Request2"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"$ref": "#/definitions/CustomRequest"
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
"required": [
|
"required": [
|
||||||
|
|
@ -2271,6 +2291,9 @@
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"$ref": "#/definitions/EmptyObject"
|
"$ref": "#/definitions/EmptyObject"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"$ref": "#/definitions/CustomResult"
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue