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

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!({
+                rmcp::handler::server::common::schema_for_empty_input()
-                    "type": "object",
-                    "properties": {}
-                }).as_object().unwrap().clone())
             })?
         }
     };

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);

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

@@ -0,0 +1,341 @@
+use std::{borrow::Cow, pin::Pin, sync::Arc};
+
+use serde::{Deserialize, Serialize};
+
+use crate::{
+    ErrorData,
+    handler::server::{
+        common::schema_for_empty_input,
+        tool::{schema_for_output, schema_for_type},
+        wrapper::{Json, Parameters},
+    },
+    model::{Icon, JsonObject, Meta, ToolAnnotations, ToolExecution},
+    schemars::JsonSchema,
+};
+
+/// Base trait to define attributes of a tool.
+///
+/// Tools implementing [`SyncTool`] or [`AsyncTool`] must implement this trait first.
+///
+/// All methods are consistent with fields of [`Tool`][crate::model::Tool].
+pub trait ToolBase {
+    /// Parameter type, will used in the invoke parameter of [`SyncTool`] or [`AsyncTool`] trait
+    ///
+    /// If the tool does not have any parameters, you **MUST** override [`input_schema`][Self::input_schema]
+    /// method. See its documentation for more details.
+    type Parameter: for<'de> Deserialize<'de> + JsonSchema + Send + Default + 'static;
+    /// Output type, will used in the invoke output of [`SyncTool`] or [`AsyncTool`] trait
+    ///
+    /// If the tool does not have any output, you **MUST** override [`output_schema`][Self::output_schema]
+    /// method. See its documentation for more details.
+    type Output: Serialize + JsonSchema + Send + 'static;
+    /// Error type, will used in the invoke output of [`SyncTool`] or [`AsyncTool`] trait
+    type Error: Into<ErrorData> + Send + 'static;
+
+    fn name() -> Cow<'static, str>;
+
+    fn title() -> Option<String> {
+        None
+    }
+    fn description() -> Option<Cow<'static, str>> {
+        None
+    }
+
+    /// Json schema for tool input.
+    ///
+    /// The default implementation generates schema based on [`Self::Parameter`] type.
+    ///
+    /// If the tool does not have any parameters, you should override this methods to return [`None`],
+    /// and when invoked, the parameter will get default values.
+    fn input_schema() -> Option<Arc<JsonObject>> {
+        Some(schema_for_type::<Parameters<Self::Parameter>>())
+    }
+
+    /// Json schema for tool output.
+    ///
+    /// The default implementation generates schema based on [`Self::Output`] type.
+    ///
+    /// If the tool does not have any output, you should override this methods to return [`None`].
+    fn output_schema() -> Option<Arc<JsonObject>> {
+        Some(schema_for_output::<Self::Output>().unwrap_or_else(|e| {
+            panic!(
+                "Invalid output schema for ToolBase::Output type `{0}`: {1}",
+                std::any::type_name::<Self::Output>(),
+                e,
+            );
+        }))
+    }
+
+    fn annotations() -> Option<ToolAnnotations> {
+        None
+    }
+    fn execution() -> Option<ToolExecution> {
+        None
+    }
+    fn icons() -> Option<Vec<Icon>> {
+        None
+    }
+    fn meta() -> Option<Meta> {
+        None
+    }
+}
+
+/// Synchronous version of a tool.
+///
+/// Consider using [`AsyncTool`] if your workflow involves asynchronous operations.
+/// Examples are shown in [the module-level documentation][crate::handler::server::router::tool].
+pub trait SyncTool<S: Sync + Send + 'static>: ToolBase {
+    fn invoke(service: &S, param: Self::Parameter) -> Result<Self::Output, Self::Error>;
+}
+
+/// Asynchronous version of a tool.
+///
+/// Consider using [`SyncTool`] if your workflow does not involve asynchronous operations.
+/// Examples are shown in [the module-level documentation][crate::handler::server::router::tool].
+pub trait AsyncTool<S: Sync + Send + 'static>: ToolBase {
+    fn invoke(
+        service: &S,
+        param: Self::Parameter,
+    ) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send;
+}
+
+pub(crate) fn tool_attribute<T: ToolBase>() -> crate::model::Tool {
+    crate::model::Tool {
+        name: T::name(),
+        title: T::title(),
+        description: T::description(),
+        input_schema: T::input_schema().unwrap_or_else(schema_for_empty_input),
+        output_schema: T::output_schema(),
+        annotations: T::annotations(),
+        execution: T::execution(),
+        icons: T::icons(),
+        meta: T::meta(),
+    }
+}
+
+pub(crate) fn sync_tool_wrapper<S: Sync + Send + 'static, T: SyncTool<S>>(
+    service: &S,
+    Parameters(params): Parameters<T::Parameter>,
+) -> Result<Json<T::Output>, ErrorData> {
+    T::invoke(service, params).map(Json).map_err(Into::into)
+}
+
+pub(crate) fn sync_tool_wrapper_with_empty_params<S: Sync + Send + 'static, T: SyncTool<S>>(
+    service: &S,
+) -> Result<Json<T::Output>, ErrorData> {
+    T::invoke(service, T::Parameter::default())
+        .map(Json)
+        .map_err(Into::into)
+}
+
+#[expect(clippy::type_complexity)]
+pub(crate) fn async_tool_wrapper<S: Sync + Send + 'static, T: AsyncTool<S>>(
+    service: &S,
+    Parameters(params): Parameters<T::Parameter>,
+) -> Pin<Box<dyn Future<Output = Result<Json<T::Output>, ErrorData>> + Send + '_>> {
+    Box::pin(async move {
+        T::invoke(service, params)
+            .await
+            .map(Json)
+            .map_err(Into::into)
+    })
+}
+
+#[expect(clippy::type_complexity)]
+pub(crate) fn async_tool_wrapper_with_empty_params<S: Sync + Send + 'static, T: AsyncTool<S>>(
+    service: &S,
+) -> Pin<Box<dyn Future<Output = Result<Json<T::Output>, ErrorData>> + Send + '_>> {
+    Box::pin(async move {
+        T::invoke(service, T::Parameter::default())
+            .await
+            .map(Json)
+            .map_err(Into::into)
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate as rmcp;
+    use crate::tool; // workaround for macros
+
+    #[derive(Deserialize, schemars::JsonSchema, Default)]
+    struct AddParameter {
+        left: usize,
+        right: usize,
+    }
+    #[derive(Serialize, schemars::JsonSchema, PartialEq, Debug)]
+    struct AddOutput {
+        sum: usize,
+    }
+
+    struct MacroBasedToolServer;
+
+    impl MacroBasedToolServer {
+        #[expect(unused)]
+        #[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),
+            })
+        }
+
+        #[expect(unused)]
+        #[tool(name = "empty", description = "Empty tool")]
+        fn empty(&self) {}
+    }
+
+    struct AddTool;
+    impl ToolBase for AddTool {
+        type Parameter = AddParameter;
+        type Output = AddOutput;
+        type Error = ErrorData;
+
+        fn name() -> Cow<'static, str> {
+            "adder".into()
+        }
+
+        fn description() -> Option<Cow<'static, str>> {
+            Some("Modular add two integers".into())
+        }
+    }
+    impl SyncTool<TraitBasedToolServer> for AddTool {
+        fn invoke(
+            _service: &TraitBasedToolServer,
+            AddParameter { left, right }: Self::Parameter,
+        ) -> Result<Self::Output, Self::Error> {
+            Ok(AddOutput {
+                sum: left.wrapping_add(right),
+            })
+        }
+    }
+    impl AsyncTool<TraitBasedToolServer> for AddTool {
+        async fn invoke(
+            _service: &TraitBasedToolServer,
+            AddParameter { left, right }: Self::Parameter,
+        ) -> Result<Self::Output, Self::Error> {
+            Ok(AddOutput {
+                sum: left.wrapping_add(right),
+            })
+        }
+    }
+
+    enum EmptyToolCustomError {
+        Internal,
+        InvalidParams,
+    }
+    impl From<EmptyToolCustomError> for ErrorData {
+        fn from(value: EmptyToolCustomError) -> Self {
+            match value {
+                EmptyToolCustomError::Internal => Self::internal_error("internal error", None),
+                EmptyToolCustomError::InvalidParams => Self::invalid_params("invalid params", None),
+            }
+        }
+    }
+
+    struct EmptyTool;
+    impl ToolBase for EmptyTool {
+        type Parameter = ();
+        type Output = ();
+        type Error = EmptyToolCustomError;
+
+        fn name() -> Cow<'static, str> {
+            "empty".into()
+        }
+
+        fn description() -> Option<Cow<'static, str>> {
+            Some("Empty tool".into())
+        }
+
+        fn input_schema() -> Option<Arc<JsonObject>> {
+            None
+        }
+
+        fn output_schema() -> Option<Arc<JsonObject>> {
+            None
+        }
+    }
+    impl SyncTool<TraitBasedToolServer> for EmptyTool {
+        fn invoke(
+            _service: &TraitBasedToolServer,
+            _param: Self::Parameter,
+        ) -> Result<Self::Output, Self::Error> {
+            Err(EmptyToolCustomError::Internal)
+        }
+    }
+    impl AsyncTool<TraitBasedToolServer> for EmptyTool {
+        async fn invoke(
+            _service: &TraitBasedToolServer,
+            _param: Self::Parameter,
+        ) -> Result<Self::Output, Self::Error> {
+            Err(EmptyToolCustomError::InvalidParams)
+        }
+    }
+
+    struct TraitBasedToolServer;
+
+    #[test]
+    fn test_macro_and_trait_have_same_attrs() {
+        let macro_attrs = MacroBasedToolServer::add_tool_attr();
+        let trait_attrs = tool_attribute::<AddTool>();
+        assert_eq!(macro_attrs, trait_attrs);
+    }
+
+    #[test]
+    fn test_macro_and_trait_have_same_attrs_for_empty_tool() {
+        let macro_attrs = MacroBasedToolServer::empty_tool_attr();
+        let trait_attrs = tool_attribute::<EmptyTool>();
+        assert_eq!(macro_attrs, trait_attrs);
+    }
+
+    #[test]
+    fn test_sync_tool_wrapper_happy_path() {
+        let left = 1;
+        let right = 2;
+        let result = sync_tool_wrapper::<_, AddTool>(
+            &TraitBasedToolServer,
+            Parameters(AddParameter { left, right }),
+        );
+        assert!(result.is_ok());
+        if let Ok(result) = result {
+            assert_eq!(result.0, AddOutput { sum: 3 });
+        }
+    }
+
+    #[tokio::test]
+    async fn test_async_tool_wrapper_happy_path() {
+        let left = 1;
+        let right = 2;
+        let result = async_tool_wrapper::<_, AddTool>(
+            &TraitBasedToolServer,
+            Parameters(AddParameter { left, right }),
+        )
+        .await;
+        assert!(result.is_ok());
+        if let Ok(result) = result {
+            assert_eq!(result.0, AddOutput { sum: 3 });
+        }
+    }
+
+    #[test]
+    fn test_sync_tool_wrapper_error_conversion() {
+        let result = sync_tool_wrapper::<_, EmptyTool>(&TraitBasedToolServer, Parameters(()));
+        assert!(result.is_err());
+        if let Err(result) = result {
+            assert_eq!(result, ErrorData::internal_error("internal error", None));
+        }
+    }
+
+    #[tokio::test]
+    async fn test_async_tool_wrapper_error_conversion() {
+        let result =
+            async_tool_wrapper::<_, EmptyTool>(&TraitBasedToolServer, Parameters(())).await;
+        assert!(result.is_err());
+        if let Err(result) = result {
+            assert_eq!(result, ErrorData::invalid_params("invalid params", None));
+        }
+    }
+}