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.
2025-12-18 09:41:06 -08:00 · 2025-12-18 09:41:06 -08:00 · e0faf1ed51
commit e0faf1ed51
parent 2e3cc4a973
5 changed files with 524 additions and 38 deletions
--- a/crates/rmcp/tests/test_custom_request.rs
+++ b/crates/rmcp/tests/test_custom_request.rs
@ -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(())
+}
--- a/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema.json
+++ b/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema.json
@ -288,6 +288,9 @@
        },
        {
          "$ref": "#/definitions/EmptyObject"
+        },
+        {
+          "$ref": "#/definitions/CustomResult"
        }
      ]
    },
@ -409,6 +412,22 @@
        "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": {
      "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": [
@ -444,7 +463,8 @@
    },
    "EmptyObject": {
      "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": {
      "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/CustomRequest"
        }
      ],
      "required": [
--- a/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema_current.json
+++ b/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema_current.json
@ -288,6 +288,9 @@
        },
        {
          "$ref": "#/definitions/EmptyObject"
+        },
+        {
+          "$ref": "#/definitions/CustomResult"
        }
      ]
    },
@ -409,6 +412,22 @@
        "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": {
      "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": [
@ -444,7 +463,8 @@
    },
    "EmptyObject": {
      "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": {
      "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/CustomRequest"
        }
      ],
      "required": [
--- a/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema.json
+++ b/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema.json
@ -392,19 +392,6 @@
        "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",
@ -606,6 +593,35 @@
        "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": {
      "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": [
@ -682,7 +698,8 @@
    },
    "EmptyObject": {
      "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": {
      "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/CustomRequest"
        }
      ],
      "required": [
@ -2271,6 +2291,9 @@
        },
        {
          "$ref": "#/definitions/EmptyObject"
+        },
+        {
+          "$ref": "#/definitions/CustomResult"
        }
      ]
    },
--- a/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema_current.json
+++ b/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema_current.json
@ -392,19 +392,6 @@
        "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",
@ -606,6 +593,35 @@
        "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": {
      "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": [
@ -682,7 +698,8 @@
    },
    "EmptyObject": {
      "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": {
      "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/CustomRequest"
        }
      ],
      "required": [
@ -2271,6 +2291,9 @@
        },
        {
          "$ref": "#/definitions/EmptyObject"
+        },
+        {
+          "$ref": "#/definitions/CustomResult"
        }
      ]
    },