feat: add trait-based tool declaration (#677)

* feat: add trait-based tool declaration

* fix: typo

* fix: add docs, make more idomatic patterns, allow for empty parameters and return types

* fix: format code

* fix: add default trait

* fix: docs typo

Author: Evian-Zhang

crates/rmcp-macros/src/tool.rs

@@ -238,10 +238,7 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
             // if not found, use a default empty JSON schema object
             // TODO: should be updated according to the new specifications
             syn::parse2::<Expr>(quote! {
-                std::sync::Arc::new(serde_json::json!({
-                    "type": "object",
-                    "properties": {}
-                }).as_object().unwrap().clone())
+                rmcp::handler::server::common::schema_for_empty_input()
             })?
         }
     };

crates/rmcp/README.md

@@ -130,7 +130,7 @@ async fn calculate(&self, params: Parameters<CalculationRequest>) -> Result<Json
 # }
 ```
 
-The `#[tool]` macro automatically generates an output schema from the `CalculationResult` type.
+The `#[tool]` macro automatically generates an output schema from the `CalculationResult` type. See the [documentation of `tool` module](crate::handler::server::router::tool) for more instructions.
 
 ## Tasks
 

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

@@ -48,6 +48,20 @@ pub fn schema_for_type<T: JsonSchema + std::any::Any>() -> Arc<JsonObject> {
     })
 }
 
+// TODO: should be updated according to the new specifications
+/// Schema used when input is empty.
+pub fn schema_for_empty_input() -> Arc<JsonObject> {
+    std::sync::Arc::new(
+        serde_json::json!({
+            "type": "object",
+            "properties": {}
+        })
+        .as_object()
+        .unwrap()
+        .clone(),
+    )
+}
+
 /// Generate and validate a JSON schema for outputSchema (must have root type "object").
 pub fn schema_for_output<T: JsonSchema + std::any::Any>() -> Result<Arc<JsonObject>, String> {
     thread_local! {

crates/rmcp/src/handler/server/router/tool.rs

@@ -1,7 +1,132 @@
+//! Tools for MCP servers.
+//!
+//! It's straightforward to define tools using [`tool_router`][crate::tool_router] and
+//! [`tool`][crate::tool] macro.
+//!
+//! ```rust
+//! # use rmcp::{
+//! #     tool_router, tool,
+//! #     handler::server::{wrapper::{Parameters, Json}, tool::ToolRouter},
+//! #     schemars
+//! # };
+//! # use serde::{Serialize, Deserialize};
+//! struct Server {
+//!     tool_router: ToolRouter<Self>,
+//! }
+//! #[derive(Deserialize, schemars::JsonSchema, Default)]
+//! struct AddParameter {
+//!     left: usize,
+//!     right: usize
+//! }
+//! #[derive(Serialize, schemars::JsonSchema)]
+//! struct AddOutput {
+//!     sum: usize
+//! }
+//! #[tool_router]
+//! impl Server {
+//!     #[tool(name = "adder", description = "Modular add two integers")]
+//!     fn add(
+//!         &self,
+//!         Parameters(AddParameter { left, right }): Parameters<AddParameter>
+//!     ) -> Json<AddOutput> {
+//!         Json(AddOutput { sum: left.wrapping_add(right) })
+//!     }
+//! }
+//! ```
+//!
+//! Using the macro-based code pattern above is suitable for small MCP servers with simple interfaces.
+//! When the business logic become larger, it is recommended that each tool should reside
+//! in individual file, combined into MCP server using [`SyncTool`] and [`AsyncTool`] traits.
+//!
+//! ```rust
+//! # use rmcp::{
+//! #     handler::server::{
+//! #         tool::ToolRouter,
+//! #         router::tool::{SyncTool, AsyncTool, ToolBase},
+//! #     },
+//! #     schemars, ErrorData
+//! # };
+//! # pub struct MyCustomError;
+//! # impl From<MyCustomError> for ErrorData {
+//! #     fn from(err: MyCustomError) -> ErrorData { unimplemented!() }
+//! # }
+//! # use serde::{Serialize, Deserialize};
+//! # use std::borrow::Cow;
+//! // In tool1.rs
+//! pub struct ComplexTool1;
+//! #[derive(Deserialize, schemars::JsonSchema, Default)]
+//! pub struct ComplexTool1Input { /* ... */ }
+//! #[derive(Serialize, schemars::JsonSchema)]
+//! pub struct ComplexTool1Output { /* ... */ }
+//!
+//! impl ToolBase for ComplexTool1 {
+//!     type Parameter = ComplexTool1Input;
+//!     type Output = ComplexTool1Output;
+//!     type Error = MyCustomError;
+//!     fn name() -> Cow<'static, str> {
+//!         "complex-tool1".into()
+//!     }
+//!
+//!     fn description() -> Option<Cow<'static, str>> {
+//!         Some("...".into())
+//!     }
+//! }
+//! impl SyncTool<MyToolServer> for ComplexTool1 {
+//!     fn invoke(service: &MyToolServer, param: Self::Parameter) -> Result<Self::Output, Self::Error> {
+//!         // ...
+//! #       unimplemented!()
+//!     }
+//! }
+//! // In tool2.rs
+//! pub struct ComplexTool2;
+//! #[derive(Deserialize, schemars::JsonSchema, Default)]
+//! pub struct ComplexTool2Input { /* ... */ }
+//! #[derive(Serialize, schemars::JsonSchema)]
+//! pub struct ComplexTool2Output { /* ... */ }
+//!
+//! impl ToolBase for ComplexTool2 {
+//!     type Parameter = ComplexTool2Input;
+//!     type Output = ComplexTool2Output;
+//!     type Error = MyCustomError;
+//!     fn name() -> Cow<'static, str> {
+//!         "complex-tool2".into()
+//!     }
+//!
+//!     fn description() -> Option<Cow<'static, str>> {
+//!         Some("...".into())
+//!     }
+//! }
+//! impl AsyncTool<MyToolServer> for ComplexTool2 {
+//!     async fn invoke(service: &MyToolServer, param: Self::Parameter) -> Result<Self::Output, Self::Error> {
+//!         // ...
+//! #       unimplemented!()
+//!     }
+//! }
+//!
+//! // In tool_router.rs
+//! struct MyToolServer {
+//!     tool_router: ToolRouter<Self>,
+//! }
+//! impl MyToolServer {
+//!     pub fn tool_router() -> ToolRouter<Self> {
+//!         ToolRouter::new()
+//!             .with_sync_tool::<ComplexTool1>()
+//!             .with_async_tool::<ComplexTool2>()
+//!     }
+//! }
+//! ```
+//!
+//! It's also possible to use macro-based and trait-based tool definition together: Since
+//! [`ToolRouter`] implements [`Add`][std::ops::Add], you can add two tool routers into final
+//! router as showed in [the documentation of `tool_router`][crate::tool_router].
+
+mod tool_traits;
+
 use std::{borrow::Cow, sync::Arc};
 
 use futures::{FutureExt, future::BoxFuture};
 use schemars::JsonSchema;
+pub use tool_traits::{AsyncTool, SyncTool, ToolBase};
 
 use crate::{
     handler::server::{
@@ -219,6 +344,42 @@ where
         self
     }
 
+    /// Add a tool that implements [`SyncTool`]
+    pub fn with_sync_tool<T>(self) -> Self
+    where
+        T: SyncTool<S> + 'static,
+    {
+        if T::input_schema().is_some() {
+            self.with_route((
+                tool_traits::tool_attribute::<T>(),
+                tool_traits::sync_tool_wrapper::<S, T>,
+            ))
+        } else {
+            self.with_route((
+                tool_traits::tool_attribute::<T>(),
+                tool_traits::sync_tool_wrapper_with_empty_params::<S, T>,
+            ))
+        }
+    }
+
+    /// Add a tool that implements [`AsyncTool`]
+    pub fn with_async_tool<T>(self) -> Self
+    where
+        T: AsyncTool<S> + 'static,
+    {
+        if T::input_schema().is_some() {
+            self.with_route((
+                tool_traits::tool_attribute::<T>(),
+                tool_traits::async_tool_wrapper::<S, T>,
+            ))
+        } else {
+            self.with_route((
+                tool_traits::tool_attribute::<T>(),
+                tool_traits::async_tool_wrapper_with_empty_params::<S, T>,
+            ))
+        }
+    }
+
     pub fn add_route(&mut self, item: ToolRoute<S>) {
         let new_name = &item.attr.name;
         validate_and_warn_tool_name(new_name);