Skip to content

Commit 1f7f4d3

Browse files
4t1454t145
authored
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
1 parent 6d31905 commit 1f7f4d3

35 files changed

+1498
-1147
lines changed

‎README.md

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -124,5 +124,10 @@ See [oauth_support](docs/OAUTH_SUPPORT.md) for details.
124124
## Related Projects
125125
- [containerd-mcp-server](https://github.com/jokemanfire/mcp-containerd) - A containerd-based MCP server implementation
126126

127-
## Development with Dev Container
128-
See [docs/DEVCONTAINER.md](docs/DEVCONTAINER.md) for instructions on using Dev Container for development.
127+
## Development
128+
129+
### Tips for Contributors
130+
See [docs/CONTRIBUTE.MD](docs/CONTRIBUTE.MD) to get some tips for contributing.
131+
132+
### Using Dev Container
133+
If you want to use dev container, see [docs/DEVCONTAINER.md](docs/DEVCONTAINER.md) for instructions on using Dev Container for development.

‎crates/rmcp-macros/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ syn = {version = "2", features = ["full"]}
1919
quote = "1"
2020
proc-macro2 = "1"
2121
serde_json = "1.0"
22-
22+
darling = { version = "0.20" }
2323

2424
[features]
2525
[dev-dependencies]

‎crates/rmcp-macros/README.md

Lines changed: 127 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -10,41 +10,149 @@ This library primarily provides the following macros:
1010

1111
## Usage
1212

13-
### Tool Macro
13+
### tool
1414

15-
Mark a function as a tool:
15+
This macro is used to mark a function as a tool handler.
1616

17-
```rust ignore
18-
#[tool]
19-
fn calculator(&self, #[tool(param)] a: i32, #[tool(param)] b: i32) -> Result<CallToolResult, Error> {
20-
// Implement tool functionality
21-
Ok(CallToolResult::success(vec![Content::text((a + b).to_string())]))
22-
}
17+
This will generate a function that return the attribute of this tool, with type `rmcp::model::Tool`.
18+
19+
#### Usage
20+
21+
| feied | type | usage |
22+
| :- | :- | :- |
23+
| `name` | `String` | The name of the tool. If not provided, it defaults to the function name. |
24+
| `description` | `String` | A description of the tool. The document of this function will be used. |
25+
| `input_schema` | `Expr` | A JSON Schema object defining the expected parameters for the tool. If not provide, if will use the json schema of its argument with type `Parameters<T>` |
26+
| `annotations` | `ToolAnnotationsAttribute` | Additional tool information. Defaults to `None`. |
2327

28+
#### Example
29+
30+
```rust
31+
#[tool(name = "my_tool", description = "This is my tool", annotations(title = "我的工具", read_only_hint = true))]
32+
pub async fn my_tool(param: Parameters<MyToolParam>) {
33+
// handling tool request
34+
}
2435
```
2536

26-
Use on an impl block to automatically register multiple tools:
37+
### tool_router
38+
39+
This macro is used to generate a tool router based on functions marked with `#[rmcp::tool]` in an implementation block.
40+
41+
It creates a function that returns a `ToolRouter` instance.
42+
43+
In most case, you need to add a field for handler to store the router information and initialize it when creating handler, or store it with a static variable.
44+
45+
#### Usage
46+
47+
| feied | type | usage |
48+
| :- | :- | :- |
49+
| `router` | `Ident` | The name of the router function to be generated. Defaults to `tool_router`. |
50+
| `vis` | `Visibility` | The visibility of the generated router function. Defaults to empty. |
2751

28-
```rust ignore
29-
#[tool(tool_box)]
30-
impl MyHandler {
52+
#### Example
53+
54+
```rust
55+
#[tool_router]
56+
impl MyToolHandler {
3157
#[tool]
32-
fn tool1(&self) -> Result<CallToolResult, Error> {
33-
// Tool 1 implementation
58+
pub fn my_tool() {
59+
3460
}
35-
36-
#[tool]
37-
fn tool2(&self) -> Result<CallToolResult, Error> {
38-
// Tool 2 implementation
61+
62+
pub fn new() -> Self {
63+
Self {
64+
// the default name of tool router will be `tool_router`
65+
tool_router: Self::tool_router(),
66+
}
3967
}
4068
}
4169
```
4270

71+
Or specify the visibility and router name, which would be helpful when you want to combine multiple routers into one:
72+
73+
```rust
74+
mod a {
75+
#[tool_router(router = tool_router_a, vis = pub)]
76+
impl MyToolHandler {
77+
#[tool]
78+
fn my_tool_a() {
79+
80+
}
81+
}
82+
}
83+
84+
mod b {
85+
#[tool_router(router = tool_router_b, vis = pub)]
86+
impl MyToolHandler {
87+
#[tool]
88+
fn my_tool_b() {
89+
90+
}
91+
}
92+
}
93+
94+
impl MyToolHandler {
95+
fn new() -> Self {
96+
Self {
97+
tool_router: self::tool_router_a() + self::tool_router_b(),
98+
}
99+
}
100+
}
101+
102+
103+
### tool_handler
104+
105+
This macro will generate the handler for `tool_call` and `list_tools` methods in the implementation block, by using an existing `ToolRouter` instance.
106+
107+
#### Usage
108+
109+
| field | type | usage |
110+
| :- | :- | :- |
111+
| `router` | `Expr` | The expression to access the `ToolRouter` instance. Defaults to `self.tool_router`. |
112+
113+
#### Example
114+
```rust
115+
#[tool_handler]
116+
impl ServerHandler for MyToolHandler {
117+
// ...implement other handler
118+
}
119+
```
120+
121+
or using a custom router expression:
122+
```rust
123+
#[tool_handler(router = self.get_router().await)]
124+
impl ServerHandler for MyToolHandler {
125+
// ...implement other handler
126+
}
127+
```
128+
129+
#### Explained
130+
This macro will be expended to something like this:
131+
```rust
132+
impl ServerHandler for MyToolHandler {
133+
async fn call_tool(
134+
&self,
135+
request: CallToolRequestParam,
136+
context: RequestContext<RoleServer>,
137+
) -> Result<CallToolResult, rmcp::Error> {
138+
let tcc = ToolCallContext::new(self, request, context);
139+
self.tool_router.call(tcc).await
140+
}
141+
142+
async fn list_tools(
143+
&self,
144+
_request: Option<PaginatedRequestParam>,
145+
_context: RequestContext<RoleServer>,
146+
) -> Result<ListToolsResult, rmcp::Error> {
147+
let items = self.tool_router.list_all();
148+
Ok(ListToolsResult::with_all_items(items))
149+
}
150+
}
151+
```
43152

44153

45154
## Advanced Features
46155

47-
- Support for parameter aggregation (`#[tool(aggr)]`)
48156
- Support for custom tool names and descriptions
49157
- Automatic generation of tool descriptions from documentation comments
50158
- JSON Schema generation for tool parameters

‎crates/rmcp-macros/src/lib.rs

Lines changed: 152 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,161 @@
22
use proc_macro::TokenStream;
33

44
mod tool;
5-
5+
mod tool_handler;
6+
mod tool_router;
7+
/// # tool
8+
///
9+
/// This macro is used to mark a function as a tool handler.
10+
///
11+
/// This will generate a function that return the attribute of this tool, with type `rmcp::model::Tool`.
12+
///
13+
/// ## Usage
14+
///
15+
/// | feied | type | usage |
16+
/// | :- | :- | :- |
17+
/// | `name` | `String` | The name of the tool. If not provided, it defaults to the function name. |
18+
/// | `description` | `String` | A description of the tool. The document of this function will be used. |
19+
/// | `input_schema` | `Expr` | A JSON Schema object defining the expected parameters for the tool. If not provide, if will use the json schema of its argument with type `Parameters<T>` |
20+
/// | `annotations` | `ToolAnnotationsAttribute` | Additional tool information. Defaults to `None`. |
21+
///
22+
/// ## Example
23+
///
24+
/// ```rust,ignore
25+
/// #[tool(name = "my_tool", description = "This is my tool", annotations(title = "我的工具", read_only_hint = true))]
26+
/// pub async fn my_tool(param: Parameters<MyToolParam>) {
27+
/// // handling tool request
28+
/// }
29+
/// ```
630
#[proc_macro_attribute]
731
pub fn tool(attr: TokenStream, input: TokenStream) -> TokenStream {
832
tool::tool(attr.into(), input.into())
933
.unwrap_or_else(|err| err.to_compile_error())
1034
.into()
1135
}
36+
37+
/// # tool_router
38+
///
39+
/// This macro is used to generate a tool router based on functions marked with `#[rmcp::tool]` in an implementation block.
40+
///
41+
/// It creates a function that returns a `ToolRouter` instance.
42+
///
43+
/// In most case, you need to add a field for handler to store the router information and initialize it when creating handler, or store it with a static variable.
44+
/// ## Usage
45+
///
46+
/// | feied | type | usage |
47+
/// | :- | :- | :- |
48+
/// | `router` | `Ident` | The name of the router function to be generated. Defaults to `tool_router`. |
49+
/// | `vis` | `Visibility` | The visibility of the generated router function. Defaults to empty. |
50+
///
51+
/// ## Example
52+
///
53+
/// ```rust,ignore
54+
/// #[tool_router]
55+
/// impl MyToolHandler {
56+
/// #[tool]
57+
/// pub fn my_tool() {
58+
///
59+
/// }
60+
///
61+
/// pub fn new() -> Self {
62+
/// Self {
63+
/// // the default name of tool router will be `tool_router`
64+
/// tool_router: Self::tool_router(),
65+
/// }
66+
/// }
67+
/// }
68+
/// ```
69+
///
70+
/// Or specify the visibility and router name, which would be helpful when you want to combine multiple routers into one:
71+
///
72+
/// ```rust,ignore
73+
/// mod a {
74+
/// #[tool_router(router = tool_router_a, vis = pub)]
75+
/// impl MyToolHandler {
76+
/// #[tool]
77+
/// fn my_tool_a() {
78+
///
79+
/// }
80+
/// }
81+
/// }
82+
///
83+
/// mod b {
84+
/// #[tool_router(router = tool_router_b, vis = pub)]
85+
/// impl MyToolHandler {
86+
/// #[tool]
87+
/// fn my_tool_b() {
88+
///
89+
/// }
90+
/// }
91+
/// }
92+
///
93+
/// impl MyToolHandler {
94+
/// fn new() -> Self {
95+
/// Self {
96+
/// tool_router: self::tool_router_a() + self::tool_router_b(),
97+
/// }
98+
/// }
99+
/// }
100+
/// ```
101+
#[proc_macro_attribute]
102+
pub fn tool_router(attr: TokenStream, input: TokenStream) -> TokenStream {
103+
tool_router::tool_router(attr.into(), input.into())
104+
.unwrap_or_else(|err| err.to_compile_error())
105+
.into()
106+
}
107+
108+
/// # tool_handler
109+
///
110+
/// This macro will generate the handler for `tool_call` and `list_tools` methods in the implementation block, by using an existing `ToolRouter` instance.
111+
///
112+
/// ## Usage
113+
///
114+
/// | field | type | usage |
115+
/// | :- | :- | :- |
116+
/// | `router` | `Expr` | The expression to access the `ToolRouter` instance. Defaults to `self.tool_router`. |
117+
/// ## Example
118+
/// ```rust,ignore
119+
/// #[tool_handler]
120+
/// impl ServerHandler for MyToolHandler {
121+
/// // ...implement other handler
122+
/// }
123+
/// ```
124+
///
125+
/// or using a custom router expression:
126+
/// ```rust,ignore
127+
/// #[tool_handler(router = self.get_router().await)]
128+
/// impl ServerHandler for MyToolHandler {
129+
/// // ...implement other handler
130+
/// }
131+
/// ```
132+
///
133+
/// ## Explain
134+
///
135+
/// This macro will be expended to something like this:
136+
/// ```rust,ignore
137+
/// impl ServerHandler for MyToolHandler {
138+
/// async fn call_tool(
139+
/// &self,
140+
/// request: CallToolRequestParam,
141+
/// context: RequestContext<RoleServer>,
142+
/// ) -> Result<CallToolResult, rmcp::Error> {
143+
/// let tcc = ToolCallContext::new(self, request, context);
144+
/// self.tool_router.call(tcc).await
145+
/// }
146+
///
147+
/// async fn list_tools(
148+
/// &self,
149+
/// _request: Option<PaginatedRequestParam>,
150+
/// _context: RequestContext<RoleServer>,
151+
/// ) -> Result<ListToolsResult, rmcp::Error> {
152+
/// let items = self.tool_router.list_all();
153+
/// Ok(ListToolsResult::with_all_items(items))
154+
/// }
155+
/// }
156+
/// ```
157+
#[proc_macro_attribute]
158+
pub fn tool_handler(attr: TokenStream, input: TokenStream) -> TokenStream {
159+
tool_handler::tool_handler(attr.into(), input.into())
160+
.unwrap_or_else(|err| err.to_compile_error())
161+
.into()
162+
}

0 commit comments

Comments
 (0)