refactor: refactor tool macros and router implementation (#261)

* refactor: refactor tool macros and router implementation

- Updated the `#[tool(tool_box)]` macro to `#[tool_router]` across various modules for consistency.
- Enhanced the `Calculator`, `Counter`, and `GenericService` structs to utilize `ToolRouter` for handling tool calls.
- Introduced `Parameters` struct for better parameter handling in tool functions.
- Added new methods for listing tools and calling tools in server handlers.
- Improved test cases to reflect changes in tool routing and parameter handling.
- Updated documentation and examples to align with the new router structure.

* fix: fix fmt and build error

* fix: fix test failure

* docs: documents for macros, fix ci

* fix: fix ci

* fix: fix wrongly replaced documents

* fix: remove useless file

* fix: change the parameter format for tool_router

* fix: update extract_doc_line to handle existing documentation and clean up unused code in server handler

* doc: update document for macro and examples

* doc: update readme and add contribute guide

* fix: fix type

crates/rmcp/tests/common/calculator.rs

@@ -1,7 +1,9 @@
+#![allow(dead_code)]
 use rmcp::{
     ServerHandler,
+    handler::server::{router::tool::ToolRouter, tool::Parameters},
     model::{ServerCapabilities, ServerInfo},
-    schemars, tool,
+    schemars, tool, tool_router,
 };
 #[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
 pub struct SumRequest {
@@ -9,30 +11,46 @@ pub struct SumRequest {
     pub a: i32,
     pub b: i32,
 }
-#[derive(Debug, Clone, Default)]
-pub struct Calculator;
-#[tool(tool_box)]
+
+#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
+pub struct SubRequest {
+    #[schemars(description = "the left hand side number")]
+    pub a: i32,
+    #[schemars(description = "the right hand side number")]
+    pub b: i32,
+}
+#[derive(Debug, Clone)]
+pub struct Calculator {
+    tool_router: ToolRouter<Self>,
+}
+
+impl Calculator {
+    pub fn new() -> Self {
+        Self {
+            tool_router: Self::tool_router(),
+        }
+    }
+}
+
+impl Default for Calculator {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[tool_router]
 impl Calculator {
     #[tool(description = "Calculate the sum of two numbers")]
-    fn sum(&self, #[tool(aggr)] SumRequest { a, b }: SumRequest) -> String {
+    fn sum(&self, Parameters(SumRequest { a, b }): Parameters<SumRequest>) -> String {
         (a + b).to_string()
     }
 
     #[tool(description = "Calculate the sub of two numbers")]
-    fn sub(
-        &self,
-        #[tool(param)]
-        #[schemars(description = "the left hand side number")]
-        a: i32,
-        #[tool(param)]
-        #[schemars(description = "the right hand side number")]
-        b: i32,
-    ) -> String {
+    fn sub(&self, Parameters(SubRequest { a, b }): Parameters<SubRequest>) -> String {
         (a - b).to_string()
     }
 }
 
-#[tool(tool_box)]
 impl ServerHandler for Calculator {
     fn get_info(&self) -> ServerInfo {
         ServerInfo {

crates/rmcp/tests/test_complex_schema.rs

@@ -1,4 +1,6 @@
-use rmcp::{Error as McpError, model::*, schemars, tool};
+use rmcp::{
+    Error as McpError, handler::server::tool::Parameters, model::*, schemars, tool, tool_router,
+};
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
@@ -24,7 +26,7 @@ pub struct ChatRequest {
 #[derive(Clone, Default)]
 pub struct Demo;
 
-#[tool(tool_box)]
+#[tool_router]
 impl Demo {
     pub fn new() -> Self {
         Self
@@ -33,9 +35,9 @@ impl Demo {
     #[tool(description = "LLM")]
     async fn chat(
         &self,
-        #[tool(aggr)] chat_request: ChatRequest,
+        chat_request: Parameters<ChatRequest>,
     ) -> Result<CallToolResult, McpError> {
-        let content = Content::json(chat_request)?;
+        let content = Content::json(chat_request.0)?;
         Ok(CallToolResult::success(vec![content]))
     }
 }

crates/rmcp/tests/test_tool_macro_annotations.rs

@@ -1,9 +1,11 @@
 #[cfg(test)]
 mod tests {
-    use rmcp::{ServerHandler, tool};
+    use rmcp::{ServerHandler, handler::server::router::tool::ToolRouter, tool, tool_handler};
 
     #[derive(Debug, Clone, Default)]
-    pub struct AnnotatedServer {}
+    pub struct AnnotatedServer {
+        tool_router: ToolRouter<AnnotatedServer>,
+    }
 
     impl AnnotatedServer {
         // Tool with inline comments for documentation
@@ -11,29 +13,14 @@ mod tests {
         /// This is used to test tool annotations
         #[tool(
             name = "direct-annotated-tool",
-            annotations = {
-                title: "Annotated Tool", 
-                readOnlyHint: true
-            }
+            annotations(title = "Annotated Tool", read_only_hint = true)
         )]
-        pub async fn direct_annotated_tool(&self, #[tool(param)] input: String) -> String {
+        pub async fn direct_annotated_tool(&self, input: String) -> String {
             format!("Direct: {}", input)
         }
     }
-
-    impl ServerHandler for AnnotatedServer {
-        async fn call_tool(
-            &self,
-            request: rmcp::model::CallToolRequestParam,
-            context: rmcp::service::RequestContext<rmcp::RoleServer>,
-        ) -> Result<rmcp::model::CallToolResult, rmcp::Error> {
-            let tcc = rmcp::handler::server::tool::ToolCallContext::new(self, request, context);
-            match tcc.name() {
-                "direct-annotated-tool" => Self::direct_annotated_tool_tool_call(tcc).await,
-                _ => Err(rmcp::Error::invalid_params("method not found", None)),
-            }
-        }
-    }
+    #[tool_handler]
+    impl ServerHandler for AnnotatedServer {}
 
     #[test]
     fn test_direct_tool_attributes() {

crates/rmcp/tests/test_tool_macros.rs

@@ -1,12 +1,12 @@
 //cargo test --test test_tool_macros --features "client server"
-
+#![allow(dead_code)]
 use std::sync::Arc;
 
 use rmcp::{
     ClientHandler, ServerHandler, ServiceExt,
-    handler::server::tool::ToolCallContext,
+    handler::server::{router::tool::ToolRouter, tool::Parameters},
     model::{CallToolRequestParam, ClientInfo},
-    tool,
+    tool, tool_handler, tool_router,
 };
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
@@ -17,37 +17,38 @@ pub struct GetWeatherRequest {
     pub date: String,
 }
 
-impl ServerHandler for Server {
-    async fn call_tool(
-        &self,
-        request: rmcp::model::CallToolRequestParam,
-        context: rmcp::service::RequestContext<rmcp::RoleServer>,
-    ) -> Result<rmcp::model::CallToolResult, rmcp::Error> {
-        let tcc = ToolCallContext::new(self, request, context);
-        match tcc.name() {
-            "get-weather" => Self::get_weather_tool_call(tcc).await,
-            _ => Err(rmcp::Error::invalid_params("method not found", None)),
-        }
+#[tool_handler(router = self.tool_router)]
+impl ServerHandler for Server {}
+
+#[derive(Debug, Clone)]
+#[allow(dead_code)]
+pub struct Server {
+    tool_router: ToolRouter<Self>,
+}
+
+impl Default for Server {
+    fn default() -> Self {
+        Self::new()
     }
 }
 
-#[derive(Debug, Clone, Default)]
-pub struct Server {}
-
+#[tool_router(router = tool_router)]
 impl Server {
+    pub fn new() -> Self {
+        Self {
+            tool_router: Self::tool_router(),
+        }
+    }
+
     /// This tool is used to get the weather of a city.
-    #[tool(name = "get-weather", description = "Get the weather of a city.", vis = )]
-    pub async fn get_weather(&self, #[tool(param)] city: String) -> String {
+    #[tool(name = "get-weather", description = "Get the weather of a city.")]
+    pub async fn get_weather(&self, city: Parameters<GetWeatherRequest>) -> String {
         drop(city);
         "rain".to_string()
     }
-    #[tool(description = "Empty Parameter")]
-    async fn empty_param(&self) {}
 
-    #[tool(description = "Optional Parameter")]
-    async fn optional_param(&self, #[tool(param)] city: Option<String>) -> String {
-        city.unwrap_or_default()
-    }
+    #[tool]
+    async fn empty_param(&self) {}
 }
 
 // define generic service trait
@@ -68,13 +69,15 @@ impl DataService for MockDataService {
 #[derive(Debug, Clone)]
 pub struct GenericServer<DS: DataService> {
     data_service: Arc<DS>,
+    tool_router: ToolRouter<Self>,
 }
 
-#[tool(tool_box)]
+#[tool_router]
 impl<DS: DataService> GenericServer<DS> {
     pub fn new(data_service: DS) -> Self {
         Self {
             data_service: Arc::new(data_service),
+            tool_router: Self::tool_router(),
         }
     }
 
@@ -83,16 +86,22 @@ impl<DS: DataService> GenericServer<DS> {
         self.data_service.get_data()
     }
 }
-#[tool(tool_box)]
+
+#[tool_handler]
 impl<DS: DataService> ServerHandler for GenericServer<DS> {}
 
 #[tokio::test]
 async fn test_tool_macros() {
-    let server = Server::default();
+    let server = Server::new();
     let _attr = Server::get_weather_tool_attr();
-    let _get_weather_call_fn = Server::get_weather_tool_call;
+    let _get_weather_tool_attr_fn = Server::get_weather_tool_attr;
     let _get_weather_fn = Server::get_weather;
-    server.get_weather("harbin".into()).await;
+    server
+        .get_weather(Parameters(GetWeatherRequest {
+            city: "Harbin".into(),
+            date: "Yesterday".into(),
+        }))
+        .await;
 }
 
 #[tokio::test]
@@ -108,14 +117,14 @@ async fn test_tool_macros_with_generics() {
     let mock_service = MockDataService;
     let server = GenericServer::new(mock_service);
     let _attr = GenericServer::<MockDataService>::get_data_tool_attr();
-    let _get_data_call_fn = GenericServer::<MockDataService>::get_data_tool_call;
+    let _get_data_call_fn = GenericServer::<MockDataService>::get_data;
     let _get_data_fn = GenericServer::<MockDataService>::get_data;
     assert_eq!(server.get_data().await, "mock data");
 }
 
 #[tokio::test]
 async fn test_tool_macros_with_optional_param() {
-    let _attr = Server::optional_param_tool_attr();
+    let _attr = Server::get_weather_tool_attr();
     // println!("{_attr:?}");
     let attr_type = _attr
         .input_schema
@@ -147,49 +156,56 @@ pub struct OptionalI64TestSchema {
 }
 
 // Dummy struct to host the test tool method
-#[derive(Debug, Clone, Default)]
-pub struct OptionalSchemaTester {}
+#[derive(Debug, Clone)]
+pub struct OptionalSchemaTester {
+    tool_router: ToolRouter<Self>,
+}
 
+impl Default for OptionalSchemaTester {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl OptionalSchemaTester {
+    pub fn new() -> Self {
+        Self {
+            tool_router: Self::tool_router(),
+        }
+    }
+}
+
+#[tool_router]
 impl OptionalSchemaTester {
     // Dummy tool function using the test schema as an aggregated parameter
     #[tool(description = "A tool to test optional schema generation")]
-    async fn test_optional_aggr(&self, #[tool(aggr)] _req: OptionalFieldTestSchema) {
+    async fn test_optional(&self, _req: Parameters<OptionalFieldTestSchema>) {
         // Implementation doesn't matter for schema testing
         // Return type changed to () to satisfy IntoCallToolResult
     }
 
     // Tool function to test optional i64 handling
     #[tool(description = "A tool to test optional i64 schema generation")]
-    async fn test_optional_i64_aggr(&self, #[tool(aggr)] req: OptionalI64TestSchema) -> String {
+    async fn test_optional_i64(
+        &self,
+        Parameters(req): Parameters<OptionalI64TestSchema>,
+    ) -> String {
         match req.count {
             Some(c) => format!("Received count: {}", c),
             None => "Received null count".to_string(),
         }
     }
 }
-
+#[tool_handler]
 // Implement ServerHandler to route tool calls for OptionalSchemaTester
-impl ServerHandler for OptionalSchemaTester {
-    async fn call_tool(
-        &self,
-        request: rmcp::model::CallToolRequestParam,
-        context: rmcp::service::RequestContext<rmcp::RoleServer>,
-    ) -> Result<rmcp::model::CallToolResult, rmcp::Error> {
-        let tcc = ToolCallContext::new(self, request, context);
-        match tcc.name() {
-            "test_optional_aggr" => Self::test_optional_aggr_tool_call(tcc).await,
-            "test_optional_i64_aggr" => Self::test_optional_i64_aggr_tool_call(tcc).await,
-            _ => Err(rmcp::Error::invalid_params("method not found", None)),
-        }
-    }
-}
+impl ServerHandler for OptionalSchemaTester {}
 
 #[test]
 fn test_optional_field_schema_generation_via_macro() {
     // tests https://github.com/modelcontextprotocol/rust-sdk/issues/135
 
     // Get the attributes generated by the #[tool] macro helper
-    let tool_attr = OptionalSchemaTester::test_optional_aggr_tool_attr();
+    let tool_attr = OptionalSchemaTester::test_optional_tool_attr();
 
     // Print the actual generated schema for debugging
     println!(
@@ -257,7 +273,7 @@ async fn test_optional_i64_field_with_null_input() -> anyhow::Result<()> {
     let (server_transport, client_transport) = tokio::io::duplex(4096);
 
     // Server setup
-    let server = OptionalSchemaTester::default();
+    let server = OptionalSchemaTester::new();
     let server_handle = tokio::spawn(async move {
         server.serve(server_transport).await?.waiting().await?;
         anyhow::Ok(())
@@ -270,7 +286,7 @@ async fn test_optional_i64_field_with_null_input() -> anyhow::Result<()> {
     // Test null case
     let result = client
         .call_tool(CallToolRequestParam {
-            name: "test_optional_i64_aggr".into(),
+            name: "test_optional_i64".into(),
             arguments: Some(
                 serde_json::json!({
                     "count": null,
@@ -298,7 +314,7 @@ async fn test_optional_i64_field_with_null_input() -> anyhow::Result<()> {
     // Test Some case
     let some_result = client
         .call_tool(CallToolRequestParam {
-            name: "test_optional_i64_aggr".into(),
+            name: "test_optional_i64".into(),
             arguments: Some(
                 serde_json::json!({
                     "count": 42,

crates/rmcp/tests/test_tool_routers.rs

@@ -0,0 +1,74 @@
+use std::collections::HashMap;
+
+use futures::future::BoxFuture;
+use rmcp::{
+    ServerHandler,
+    handler::server::{
+        router::tool::ToolRouter,
+        tool::{CallToolHandler, Parameters},
+    },
+};
+
+#[derive(Debug, Default)]
+pub struct TestHandler<T: 'static = ()> {
+    pub _marker: std::marker::PhantomData<fn(*const T)>,
+}
+
+impl<T: 'static> ServerHandler for TestHandler<T> {}
+#[derive(Debug, schemars::JsonSchema, serde::Deserialize, serde::Serialize)]
+pub struct Request {
+    pub fields: HashMap<String, String>,
+}
+
+#[derive(Debug, schemars::JsonSchema, serde::Deserialize, serde::Serialize)]
+pub struct Sum {
+    pub a: i32,
+    pub b: i32,
+}
+
+#[rmcp::tool_router(router = test_router_1)]
+impl<T> TestHandler<T> {
+    #[rmcp::tool]
+    async fn async_method(&self, Parameters(Request { fields }): Parameters<Request>) {
+        drop(fields)
+    }
+}
+
+#[rmcp::tool_router(router = test_router_2)]
+impl<T> TestHandler<T> {
+    #[rmcp::tool]
+    fn sync_method(&self, Parameters(Request { fields }): Parameters<Request>) {
+        drop(fields)
+    }
+}
+
+#[rmcp::tool]
+async fn async_function<T>(
+    _callee: &TestHandler<T>,
+    Parameters(Request { fields }): Parameters<Request>,
+) {
+    drop(fields)
+}
+
+#[rmcp::tool]
+fn async_function2<T>(_callee: &TestHandler<T>) -> BoxFuture<'_, ()> {
+    Box::pin(async move {})
+}
+
+#[test]
+fn test_tool_router() {
+    let test_tool_router: ToolRouter<TestHandler<()>> = ToolRouter::<TestHandler<()>>::new()
+        .with_route((async_function_tool_attr(), async_function))
+        .with_route((async_function2_tool_attr(), async_function2))
+        + TestHandler::<()>::test_router_1()
+        + TestHandler::<()>::test_router_2();
+    let tools = test_tool_router.list_all();
+    assert_eq!(tools.len(), 4);
+    assert_handler(TestHandler::<()>::async_method);
+}
+
+fn assert_handler<S, H, A>(_handler: H)
+where
+    H: CallToolHandler<S, A>,
+{
+}

crates/rmcp/tests/test_with_js.rs

@@ -96,7 +96,7 @@ async fn test_with_js_streamable_http_client() -> anyhow::Result<()> {
 
     let service: StreamableHttpService<Calculator, LocalSessionManager> =
         StreamableHttpService::new(
-            || Ok(Calculator),
+            || Ok(Calculator::new()),
             Default::default(),
             StreamableHttpServerConfig {
                 stateful_mode: true,