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
This commit is contained in:
EvianZhang 2026-02-26 01:23:37 +08:00 • committed by GitHub
commit 6c336a90c1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

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