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
This commit is contained in:
4t145
GitHub
parent
6d3190504c
commit
1f7f4d3055
4 changed files with 1498 additions and 1147 deletions
|
|
@ -2,10 +2,161 @@
|
|||
use proc_macro::TokenStream;
|
||||
|
||||
mod tool;
|
||||
|
||||
mod tool_handler;
|
||||
mod tool_router;
|
||||
/// # tool
|
||||
///
|
||||
/// This macro is used to mark a function as a tool handler.
|
||||
///
|
||||
/// This will generate a function that return the attribute of this tool, with type `rmcp::model::Tool`.
|
||||
///
|
||||
/// ## Usage
|
||||
///
|
||||
/// | feied | type | usage |
|
||||
/// | :- | :- | :- |
|
||||
/// | `name` | `String` | The name of the tool. If not provided, it defaults to the function name. |
|
||||
/// | `description` | `String` | A description of the tool. The document of this function will be used. |
|
||||
/// | `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>` |
|
||||
/// | `annotations` | `ToolAnnotationsAttribute` | Additional tool information. Defaults to `None`. |
|
||||
///
|
||||
/// ## Example
|
||||
///
|
||||
/// ```rust,ignore
|
||||
/// #[tool(name = "my_tool", description = "This is my tool", annotations(title = "我的工具", read_only_hint = true))]
|
||||
/// pub async fn my_tool(param: Parameters<MyToolParam>) {
|
||||
/// // handling tool request
|
||||
/// }
|
||||
/// ```
|
||||
#[proc_macro_attribute]
|
||||
pub fn tool(attr: TokenStream, input: TokenStream) -> TokenStream {
|
||||
tool::tool(attr.into(), input.into())
|
||||
.unwrap_or_else(|err| err.to_compile_error())
|
||||
.into()
|
||||
}
|
||||
|
||||
/// # tool_router
|
||||
///
|
||||
/// This macro is used to generate a tool router based on functions marked with `#[rmcp::tool]` in an implementation block.
|
||||
///
|
||||
/// It creates a function that returns a `ToolRouter` instance.
|
||||
///
|
||||
/// 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.
|
||||
/// ## Usage
|
||||
///
|
||||
/// | feied | type | usage |
|
||||
/// | :- | :- | :- |
|
||||
/// | `router` | `Ident` | The name of the router function to be generated. Defaults to `tool_router`. |
|
||||
/// | `vis` | `Visibility` | The visibility of the generated router function. Defaults to empty. |
|
||||
///
|
||||
/// ## Example
|
||||
///
|
||||
/// ```rust,ignore
|
||||
/// #[tool_router]
|
||||
/// impl MyToolHandler {
|
||||
/// #[tool]
|
||||
/// pub fn my_tool() {
|
||||
///
|
||||
/// }
|
||||
///
|
||||
/// pub fn new() -> Self {
|
||||
/// Self {
|
||||
/// // the default name of tool router will be `tool_router`
|
||||
/// tool_router: Self::tool_router(),
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// Or specify the visibility and router name, which would be helpful when you want to combine multiple routers into one:
|
||||
///
|
||||
/// ```rust,ignore
|
||||
/// mod a {
|
||||
/// #[tool_router(router = tool_router_a, vis = pub)]
|
||||
/// impl MyToolHandler {
|
||||
/// #[tool]
|
||||
/// fn my_tool_a() {
|
||||
///
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// mod b {
|
||||
/// #[tool_router(router = tool_router_b, vis = pub)]
|
||||
/// impl MyToolHandler {
|
||||
/// #[tool]
|
||||
/// fn my_tool_b() {
|
||||
///
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// impl MyToolHandler {
|
||||
/// fn new() -> Self {
|
||||
/// Self {
|
||||
/// tool_router: self::tool_router_a() + self::tool_router_b(),
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
#[proc_macro_attribute]
|
||||
pub fn tool_router(attr: TokenStream, input: TokenStream) -> TokenStream {
|
||||
tool_router::tool_router(attr.into(), input.into())
|
||||
.unwrap_or_else(|err| err.to_compile_error())
|
||||
.into()
|
||||
}
|
||||
|
||||
/// # tool_handler
|
||||
///
|
||||
/// This macro will generate the handler for `tool_call` and `list_tools` methods in the implementation block, by using an existing `ToolRouter` instance.
|
||||
///
|
||||
/// ## Usage
|
||||
///
|
||||
/// | field | type | usage |
|
||||
/// | :- | :- | :- |
|
||||
/// | `router` | `Expr` | The expression to access the `ToolRouter` instance. Defaults to `self.tool_router`. |
|
||||
/// ## Example
|
||||
/// ```rust,ignore
|
||||
/// #[tool_handler]
|
||||
/// impl ServerHandler for MyToolHandler {
|
||||
/// // ...implement other handler
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// or using a custom router expression:
|
||||
/// ```rust,ignore
|
||||
/// #[tool_handler(router = self.get_router().await)]
|
||||
/// impl ServerHandler for MyToolHandler {
|
||||
/// // ...implement other handler
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// ## Explain
|
||||
///
|
||||
/// This macro will be expended to something like this:
|
||||
/// ```rust,ignore
|
||||
/// impl ServerHandler for MyToolHandler {
|
||||
/// async fn call_tool(
|
||||
/// &self,
|
||||
/// request: CallToolRequestParam,
|
||||
/// context: RequestContext<RoleServer>,
|
||||
/// ) -> Result<CallToolResult, rmcp::Error> {
|
||||
/// let tcc = ToolCallContext::new(self, request, context);
|
||||
/// self.tool_router.call(tcc).await
|
||||
/// }
|
||||
///
|
||||
/// async fn list_tools(
|
||||
/// &self,
|
||||
/// _request: Option<PaginatedRequestParam>,
|
||||
/// _context: RequestContext<RoleServer>,
|
||||
/// ) -> Result<ListToolsResult, rmcp::Error> {
|
||||
/// let items = self.tool_router.list_all();
|
||||
/// Ok(ListToolsResult::with_all_items(items))
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
#[proc_macro_attribute]
|
||||
pub fn tool_handler(attr: TokenStream, input: TokenStream) -> TokenStream {
|
||||
tool_handler::tool_handler(attr.into(), input.into())
|
||||
.unwrap_or_else(|err| err.to_compile_error())
|
||||
.into()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,348 +1,96 @@
|
|||
use std::collections::HashSet;
|
||||
|
||||
use darling::{FromMeta, ast::NestedMeta};
|
||||
use proc_macro2::TokenStream;
|
||||
use quote::{ToTokens, quote};
|
||||
use serde_json::json;
|
||||
use syn::{
|
||||
Expr, FnArg, Ident, ItemFn, ItemImpl, Lit, MetaList, PatType, Token, Type, Visibility,
|
||||
parse::Parse, parse_quote, spanned::Spanned,
|
||||
};
|
||||
|
||||
/// Stores tool annotation attributes
|
||||
#[derive(Default, Clone)]
|
||||
struct ToolAnnotationAttrs(pub serde_json::Map<String, serde_json::Value>);
|
||||
|
||||
impl Parse for ToolAnnotationAttrs {
|
||||
fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
|
||||
let mut attrs = serde_json::Map::new();
|
||||
|
||||
while !input.is_empty() {
|
||||
let key: Ident = input.parse()?;
|
||||
input.parse::<Token![:]>()?;
|
||||
let value: Lit = input.parse()?;
|
||||
let value = match value {
|
||||
Lit::Str(s) => json!(s.value()),
|
||||
Lit::Bool(b) => json!(b.value),
|
||||
_ => {
|
||||
return Err(syn::Error::new(
|
||||
key.span(),
|
||||
"annotations must be string or boolean literals",
|
||||
));
|
||||
}
|
||||
};
|
||||
attrs.insert(key.to_string(), value);
|
||||
if input.is_empty() {
|
||||
break;
|
||||
}
|
||||
input.parse::<Token![,]>()?;
|
||||
}
|
||||
|
||||
Ok(ToolAnnotationAttrs(attrs))
|
||||
}
|
||||
use quote::{ToTokens, format_ident, quote};
|
||||
use syn::{Expr, Ident, ImplItemFn, ReturnType};
|
||||
#[derive(FromMeta, Default, Debug)]
|
||||
#[darling(default)]
|
||||
pub struct ToolAttribute {
|
||||
/// The name of the tool
|
||||
pub name: Option<String>,
|
||||
pub description: Option<String>,
|
||||
/// A JSON Schema object defining the expected parameters for the tool
|
||||
pub input_schema: Option<Expr>,
|
||||
/// Optional additional tool information.
|
||||
pub annotations: Option<ToolAnnotationsAttribute>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ToolImplItemAttrs {
|
||||
tool_box: Option<Option<Ident>>,
|
||||
pub struct ResolvedToolAttribute {
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub input_schema: Expr,
|
||||
pub annotations: Expr,
|
||||
}
|
||||
|
||||
impl Parse for ToolImplItemAttrs {
|
||||
fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
|
||||
let mut tool_box = None;
|
||||
while !input.is_empty() {
|
||||
let key: Ident = input.parse()?;
|
||||
match key.to_string().as_str() {
|
||||
"tool_box" => {
|
||||
tool_box = Some(None);
|
||||
if input.lookahead1().peek(Token![=]) {
|
||||
input.parse::<Token![=]>()?;
|
||||
let value: Ident = input.parse()?;
|
||||
tool_box = Some(Some(value));
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
return Err(syn::Error::new(key.span(), "unknown attribute"));
|
||||
}
|
||||
}
|
||||
if input.is_empty() {
|
||||
break;
|
||||
}
|
||||
input.parse::<Token![,]>()?;
|
||||
}
|
||||
|
||||
Ok(ToolImplItemAttrs { tool_box })
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ToolFnItemAttrs {
|
||||
name: Option<Expr>,
|
||||
description: Option<Expr>,
|
||||
vis: Option<Visibility>,
|
||||
annotations: Option<ToolAnnotationAttrs>,
|
||||
}
|
||||
|
||||
impl Parse for ToolFnItemAttrs {
|
||||
fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
|
||||
let mut name = None;
|
||||
let mut description = None;
|
||||
let mut vis = None;
|
||||
let mut annotations = None;
|
||||
|
||||
while !input.is_empty() {
|
||||
let key: Ident = input.parse()?;
|
||||
input.parse::<Token![=]>()?;
|
||||
match key.to_string().as_str() {
|
||||
"name" => {
|
||||
let value: Expr = input.parse()?;
|
||||
name = Some(value);
|
||||
}
|
||||
"description" => {
|
||||
let value: Expr = input.parse()?;
|
||||
description = Some(value);
|
||||
}
|
||||
"vis" => {
|
||||
let value: Visibility = input.parse()?;
|
||||
vis = Some(value);
|
||||
}
|
||||
"annotations" => {
|
||||
// Parse the annotations as a nested structure
|
||||
let content;
|
||||
syn::braced!(content in input);
|
||||
let value = content.parse()?;
|
||||
annotations = Some(value);
|
||||
}
|
||||
_ => {
|
||||
return Err(syn::Error::new(key.span(), "unknown attribute"));
|
||||
}
|
||||
}
|
||||
if input.is_empty() {
|
||||
break;
|
||||
}
|
||||
input.parse::<Token![,]>()?;
|
||||
}
|
||||
|
||||
Ok(ToolFnItemAttrs {
|
||||
impl ResolvedToolAttribute {
|
||||
pub fn into_fn(self, fn_ident: Ident) -> syn::Result<ImplItemFn> {
|
||||
let Self {
|
||||
name,
|
||||
description,
|
||||
vis,
|
||||
input_schema,
|
||||
annotations,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct ToolFnParamAttrs {
|
||||
serde_meta: Vec<MetaList>,
|
||||
schemars_meta: Vec<MetaList>,
|
||||
ident: Ident,
|
||||
rust_type: Box<Type>,
|
||||
}
|
||||
|
||||
impl ToTokens for ToolFnParamAttrs {
|
||||
fn to_tokens(&self, tokens: &mut TokenStream) {
|
||||
let ident = &self.ident;
|
||||
let rust_type = &self.rust_type;
|
||||
let serde_meta = &self.serde_meta;
|
||||
let schemars_meta = &self.schemars_meta;
|
||||
tokens.extend(quote! {
|
||||
#(#[#serde_meta])*
|
||||
#(#[#schemars_meta])*
|
||||
pub #ident: #rust_type,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
|
||||
enum ToolParams {
|
||||
Aggregated {
|
||||
rust_type: PatType,
|
||||
},
|
||||
Params {
|
||||
attrs: Vec<ToolFnParamAttrs>,
|
||||
},
|
||||
#[default]
|
||||
NoParam,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ToolAttrs {
|
||||
fn_item: ToolFnItemAttrs,
|
||||
params: ToolParams,
|
||||
}
|
||||
const TOOL_IDENT: &str = "tool";
|
||||
const SERDE_IDENT: &str = "serde";
|
||||
const SCHEMARS_IDENT: &str = "schemars";
|
||||
const PARAM_IDENT: &str = "param";
|
||||
const AGGREGATED_IDENT: &str = "aggr";
|
||||
const REQ_IDENT: &str = "req";
|
||||
|
||||
pub enum ParamMarker {
|
||||
Param,
|
||||
Aggregated,
|
||||
}
|
||||
|
||||
impl Parse for ParamMarker {
|
||||
fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
|
||||
let ident: Ident = input.parse()?;
|
||||
match ident.to_string().as_str() {
|
||||
PARAM_IDENT => Ok(ParamMarker::Param),
|
||||
AGGREGATED_IDENT | REQ_IDENT => Ok(ParamMarker::Aggregated),
|
||||
_ => Err(syn::Error::new(ident.span(), "unknown attribute")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub enum ToolItem {
|
||||
Fn(ItemFn),
|
||||
Impl(ItemImpl),
|
||||
}
|
||||
|
||||
impl Parse for ToolItem {
|
||||
fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
|
||||
let lookahead = input.lookahead1();
|
||||
if lookahead.peek(Token![impl]) {
|
||||
let item = input.parse::<ItemImpl>()?;
|
||||
Ok(ToolItem::Impl(item))
|
||||
} = self;
|
||||
let description = if let Some(description) = description {
|
||||
quote! { Some(#description.into()) }
|
||||
} else {
|
||||
let item = input.parse::<ItemFn>()?;
|
||||
Ok(ToolItem::Fn(item))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// dispatch impl function item and impl block item
|
||||
pub(crate) fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
|
||||
let tool_item = syn::parse2::<ToolItem>(input)?;
|
||||
match tool_item {
|
||||
ToolItem::Fn(item) => tool_fn_item(attr, item),
|
||||
ToolItem::Impl(item) => tool_impl_item(attr, item),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn tool_impl_item(attr: TokenStream, mut input: ItemImpl) -> syn::Result<TokenStream> {
|
||||
let tool_impl_attr: ToolImplItemAttrs = syn::parse2(attr)?;
|
||||
let tool_box_ident = tool_impl_attr.tool_box;
|
||||
|
||||
// get all tool function ident
|
||||
let mut tool_fn_idents = Vec::new();
|
||||
for item in &input.items {
|
||||
if let syn::ImplItem::Fn(method) = item {
|
||||
for attr in &method.attrs {
|
||||
if attr.path().is_ident(TOOL_IDENT) {
|
||||
tool_fn_idents.push(method.sig.ident.clone());
|
||||
quote! { None }
|
||||
};
|
||||
let tokens = quote! {
|
||||
pub fn #fn_ident() -> rmcp::model::Tool {
|
||||
rmcp::model::Tool {
|
||||
name: #name.into(),
|
||||
description: #description,
|
||||
input_schema: #input_schema,
|
||||
annotations: #annotations,
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
syn::parse2::<ImplItemFn>(tokens)
|
||||
}
|
||||
}
|
||||
|
||||
// handle different cases
|
||||
if input.trait_.is_some() {
|
||||
if let Some(ident) = tool_box_ident {
|
||||
// check if there are generic parameters
|
||||
if !input.generics.params.is_empty() {
|
||||
// for trait implementation with generic parameters, directly use the already generated *_inner method
|
||||
#[derive(FromMeta, Debug, Default)]
|
||||
#[darling(default)]
|
||||
pub struct ToolAnnotationsAttribute {
|
||||
/// A human-readable title for the tool.
|
||||
pub title: Option<String>,
|
||||
|
||||
// generate call_tool method
|
||||
input.items.push(parse_quote! {
|
||||
async fn call_tool(
|
||||
&self,
|
||||
request: rmcp::model::CallToolRequestParam,
|
||||
context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
||||
) -> Result<rmcp::model::CallToolResult, rmcp::Error> {
|
||||
self.call_tool_inner(request, context).await
|
||||
}
|
||||
});
|
||||
/// If true, the tool does not modify its environment.
|
||||
///
|
||||
/// Default: false
|
||||
pub read_only_hint: Option<bool>,
|
||||
|
||||
// generate list_tools method
|
||||
input.items.push(parse_quote! {
|
||||
async fn list_tools(
|
||||
&self,
|
||||
request: Option<rmcp::model::PaginatedRequestParam>,
|
||||
context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
||||
) -> Result<rmcp::model::ListToolsResult, rmcp::Error> {
|
||||
self.list_tools_inner(request, context).await
|
||||
}
|
||||
});
|
||||
} else {
|
||||
// if there are no generic parameters, add tool box derive
|
||||
input.items.push(parse_quote!(
|
||||
rmcp::tool_box!(@derive #ident);
|
||||
));
|
||||
}
|
||||
} else {
|
||||
return Err(syn::Error::new(
|
||||
proc_macro2::Span::call_site(),
|
||||
"tool_box attribute is required for trait implementation",
|
||||
));
|
||||
}
|
||||
} else if let Some(ident) = tool_box_ident {
|
||||
// if it is a normal impl block
|
||||
if !input.generics.params.is_empty() {
|
||||
// if there are generic parameters, not use tool_box! macro, but generate code directly
|
||||
/// If true, the tool may perform destructive updates to its environment.
|
||||
/// If false, the tool performs only additive updates.
|
||||
///
|
||||
/// (This property is meaningful only when `readOnlyHint == false`)
|
||||
///
|
||||
/// Default: true
|
||||
/// A human-readable description of the tool's purpose.
|
||||
pub destructive_hint: Option<bool>,
|
||||
|
||||
// create call code for each tool function
|
||||
let match_arms = tool_fn_idents.iter().map(|ident| {
|
||||
let attr_fn = Ident::new(&format!("{}_tool_attr", ident), ident.span());
|
||||
let call_fn = Ident::new(&format!("{}_tool_call", ident), ident.span());
|
||||
quote! {
|
||||
name if name == Self::#attr_fn().name => {
|
||||
Self::#call_fn(tcc).await
|
||||
}
|
||||
}
|
||||
});
|
||||
/// If true, calling the tool repeatedly with the same arguments
|
||||
/// will have no additional effect on the its environment.
|
||||
///
|
||||
/// (This property is meaningful only when `readOnlyHint == false`)
|
||||
///
|
||||
/// Default: false.
|
||||
pub idempotent_hint: Option<bool>,
|
||||
|
||||
let tool_attrs = tool_fn_idents.iter().map(|ident| {
|
||||
let attr_fn = Ident::new(&format!("{}_tool_attr", ident), ident.span());
|
||||
quote! { Self::#attr_fn() }
|
||||
});
|
||||
/// If true, this tool may interact with an "open world" of external
|
||||
/// entities. If false, the tool's domain of interaction is closed.
|
||||
/// For example, the world of a web search tool is open, whereas that
|
||||
/// of a memory tool is not.
|
||||
///
|
||||
/// Default: true
|
||||
pub open_world_hint: Option<bool>,
|
||||
}
|
||||
|
||||
// implement call_tool method
|
||||
input.items.push(parse_quote! {
|
||||
async fn call_tool_inner(
|
||||
&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() {
|
||||
#(#match_arms,)*
|
||||
_ => Err(rmcp::Error::invalid_params("tool not found", None)),
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// implement list_tools method
|
||||
input.items.push(parse_quote! {
|
||||
async fn list_tools_inner(
|
||||
&self,
|
||||
_: Option<rmcp::model::PaginatedRequestParam>,
|
||||
_: rmcp::service::RequestContext<rmcp::RoleServer>,
|
||||
) -> Result<rmcp::model::ListToolsResult, rmcp::Error> {
|
||||
Ok(rmcp::model::ListToolsResult {
|
||||
next_cursor: None,
|
||||
tools: vec![#(#tool_attrs),*],
|
||||
})
|
||||
}
|
||||
});
|
||||
} else {
|
||||
// if there are no generic parameters, use the original tool_box! macro
|
||||
let this_type_ident = &input.self_ty;
|
||||
input.items.push(parse_quote!(
|
||||
rmcp::tool_box!(#this_type_ident {
|
||||
#(#tool_fn_idents),*
|
||||
} #ident);
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(quote! {
|
||||
#input
|
||||
})
|
||||
fn none_expr() -> Expr {
|
||||
syn::parse2::<Expr>(quote! { None }).unwrap()
|
||||
}
|
||||
|
||||
// extract doc line from attribute
|
||||
fn extract_doc_line(attr: &syn::Attribute) -> Option<String> {
|
||||
fn extract_doc_line(existing_docs: Option<String>, attr: &syn::Attribute) -> Option<String> {
|
||||
if !attr.path().is_ident("doc") {
|
||||
return None;
|
||||
}
|
||||
|
|
@ -360,391 +108,147 @@ fn extract_doc_line(attr: &syn::Attribute) -> Option<String> {
|
|||
};
|
||||
|
||||
let content = lit_str.value().trim().to_string();
|
||||
|
||||
(!content.is_empty()).then_some(content)
|
||||
match (existing_docs, content) {
|
||||
(Some(mut existing_docs), content) if !content.is_empty() => {
|
||||
existing_docs.push('\n');
|
||||
existing_docs.push_str(&content);
|
||||
Some(existing_docs)
|
||||
}
|
||||
(Some(existing_docs), _) => Some(existing_docs),
|
||||
(None, content) if !content.is_empty() => Some(content),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn tool_fn_item(attr: TokenStream, mut input_fn: ItemFn) -> syn::Result<TokenStream> {
|
||||
let mut tool_macro_attrs = ToolAttrs::default();
|
||||
let args: ToolFnItemAttrs = syn::parse2(attr)?;
|
||||
tool_macro_attrs.fn_item = args;
|
||||
// let mut fommated_fn_args: Punctuated<FnArg, Comma> = Punctuated::new();
|
||||
let mut unextractable_args_indexes = HashSet::new();
|
||||
for (index, mut fn_arg) in input_fn.sig.inputs.iter_mut().enumerate() {
|
||||
enum Caught {
|
||||
Param(ToolFnParamAttrs),
|
||||
Aggregated(PatType),
|
||||
}
|
||||
let mut caught = None;
|
||||
match &mut fn_arg {
|
||||
FnArg::Receiver(_) => {
|
||||
continue;
|
||||
}
|
||||
FnArg::Typed(pat_type) => {
|
||||
let mut serde_metas = Vec::new();
|
||||
let mut schemars_metas = Vec::new();
|
||||
let mut arg_ident = match pat_type.pat.as_ref() {
|
||||
syn::Pat::Ident(pat_ident) => Some(pat_ident.ident.clone()),
|
||||
_ => None,
|
||||
};
|
||||
let raw_attrs: Vec<_> = pat_type.attrs.drain(..).collect();
|
||||
for attr in raw_attrs {
|
||||
match &attr.meta {
|
||||
syn::Meta::List(meta_list) => {
|
||||
if meta_list.path.is_ident(TOOL_IDENT) {
|
||||
let pat_type = pat_type.clone();
|
||||
let marker = meta_list.parse_args::<ParamMarker>()?;
|
||||
match marker {
|
||||
ParamMarker::Param => {
|
||||
let Some(arg_ident) = arg_ident.take() else {
|
||||
return Err(syn::Error::new(
|
||||
proc_macro2::Span::call_site(),
|
||||
"input param must have an ident as name",
|
||||
));
|
||||
};
|
||||
caught.replace(Caught::Param(ToolFnParamAttrs {
|
||||
serde_meta: Vec::new(),
|
||||
schemars_meta: Vec::new(),
|
||||
ident: arg_ident,
|
||||
rust_type: pat_type.ty.clone(),
|
||||
}));
|
||||
}
|
||||
ParamMarker::Aggregated => {
|
||||
caught.replace(Caught::Aggregated(pat_type.clone()));
|
||||
}
|
||||
}
|
||||
} else if meta_list.path.is_ident(SERDE_IDENT) {
|
||||
serde_metas.push(meta_list.clone());
|
||||
} else if meta_list.path.is_ident(SCHEMARS_IDENT) {
|
||||
schemars_metas.push(meta_list.clone());
|
||||
} else {
|
||||
pat_type.attrs.push(attr);
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
pat_type.attrs.push(attr);
|
||||
}
|
||||
}
|
||||
}
|
||||
match caught {
|
||||
Some(Caught::Param(mut param)) => {
|
||||
param.serde_meta = serde_metas;
|
||||
param.schemars_meta = schemars_metas;
|
||||
match &mut tool_macro_attrs.params {
|
||||
ToolParams::Params { attrs } => {
|
||||
attrs.push(param);
|
||||
}
|
||||
_ => {
|
||||
tool_macro_attrs.params = ToolParams::Params { attrs: vec![param] };
|
||||
}
|
||||
}
|
||||
unextractable_args_indexes.insert(index);
|
||||
}
|
||||
Some(Caught::Aggregated(rust_type)) => {
|
||||
if let ToolParams::Params { .. } = tool_macro_attrs.params {
|
||||
return Err(syn::Error::new(
|
||||
rust_type.span(),
|
||||
"cannot mix aggregated and individual parameters",
|
||||
));
|
||||
}
|
||||
tool_macro_attrs.params = ToolParams::Aggregated { rust_type };
|
||||
unextractable_args_indexes.insert(index);
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// input_fn.sig.inputs = fommated_fn_args;
|
||||
let name = if let Some(expr) = tool_macro_attrs.fn_item.name {
|
||||
expr
|
||||
pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
|
||||
let attribute = if attr.is_empty() {
|
||||
Default::default()
|
||||
} else {
|
||||
let fn_name = &input_fn.sig.ident;
|
||||
parse_quote! {
|
||||
stringify!(#fn_name)
|
||||
}
|
||||
let attr_args = NestedMeta::parse_meta_list(attr)?;
|
||||
ToolAttribute::from_list(&attr_args)?
|
||||
};
|
||||
let tool_attr_fn_ident = Ident::new(
|
||||
&format!("{}_tool_attr", input_fn.sig.ident),
|
||||
proc_macro2::Span::call_site(),
|
||||
);
|
||||
let mut fn_item = syn::parse2::<ImplItemFn>(input.clone())?;
|
||||
let fn_ident = &fn_item.sig.ident;
|
||||
|
||||
// generate get tool attr function
|
||||
let tool_attr_fn = {
|
||||
let description = if let Some(expr) = tool_macro_attrs.fn_item.description {
|
||||
// Use explicitly provided description if available
|
||||
expr
|
||||
} else {
|
||||
// Try to extract documentation comments
|
||||
let doc_content = input_fn
|
||||
.attrs
|
||||
.iter()
|
||||
.filter_map(extract_doc_line)
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
parse_quote! {
|
||||
#doc_content.trim().to_string()
|
||||
}
|
||||
};
|
||||
let schema = match &tool_macro_attrs.params {
|
||||
ToolParams::Aggregated { rust_type } => {
|
||||
let ty = &rust_type.ty;
|
||||
let schema = quote! {
|
||||
rmcp::handler::server::tool::cached_schema_for_type::<#ty>()
|
||||
};
|
||||
schema
|
||||
}
|
||||
ToolParams::Params { attrs, .. } => {
|
||||
let (param_type, temp_param_type_name) =
|
||||
create_request_type(attrs, input_fn.sig.ident.to_string());
|
||||
let schema = quote! {
|
||||
let tool_attr_fn_ident = format_ident!("{}_tool_attr", fn_ident);
|
||||
let input_schema_expr = if let Some(input_schema) = attribute.input_schema {
|
||||
input_schema
|
||||
} else {
|
||||
// try to find some parameters wrapper in the function
|
||||
let params_ty = fn_item.sig.inputs.iter().find_map(|input| {
|
||||
if let syn::FnArg::Typed(pat_type) = input {
|
||||
if let syn::Type::Path(type_path) = &*pat_type.ty {
|
||||
if type_path
|
||||
.path
|
||||
.segments
|
||||
.last()
|
||||
.is_some_and(|type_name| type_name.ident == "Parameters")
|
||||
{
|
||||
#param_type
|
||||
rmcp::handler::server::tool::cached_schema_for_type::<#temp_param_type_name>()
|
||||
return Some(pat_type.ty.clone());
|
||||
}
|
||||
};
|
||||
schema
|
||||
}
|
||||
ToolParams::NoParam => {
|
||||
quote! {
|
||||
rmcp::handler::server::tool::cached_schema_for_type::<rmcp::model::EmptyObject>()
|
||||
}
|
||||
}
|
||||
};
|
||||
let input_fn_attrs = &input_fn.attrs;
|
||||
let input_fn_vis = &input_fn.vis;
|
||||
|
||||
let annotations_code = if let Some(annotations) = &tool_macro_attrs.fn_item.annotations {
|
||||
let annotations =
|
||||
serde_json::to_string(&annotations.0).expect("failed to serialize annotations");
|
||||
quote! {
|
||||
Some(serde_json::from_str::<rmcp::model::ToolAnnotations>(&#annotations).expect("Could not parse tool annotations"))
|
||||
}
|
||||
None
|
||||
});
|
||||
if let Some(params_ty) = params_ty {
|
||||
// if found, use the Parameters schema
|
||||
syn::parse2::<Expr>(quote! {
|
||||
rmcp::handler::server::tool::cached_schema_for_type::<#params_ty>()
|
||||
})?
|
||||
} else {
|
||||
quote! { None }
|
||||
};
|
||||
|
||||
quote! {
|
||||
#(#input_fn_attrs)*
|
||||
#input_fn_vis fn #tool_attr_fn_ident() -> rmcp::model::Tool {
|
||||
rmcp::model::Tool {
|
||||
name: #name.into(),
|
||||
description: Some(#description.into()),
|
||||
input_schema: #schema.into(),
|
||||
annotations: #annotations_code,
|
||||
}
|
||||
}
|
||||
// if not found, use the default EmptyObject schema
|
||||
syn::parse2::<Expr>(quote! {
|
||||
rmcp::handler::server::tool::cached_schema_for_type::<rmcp::model::EmptyObject>()
|
||||
})?
|
||||
}
|
||||
};
|
||||
|
||||
// generate wrapped tool function
|
||||
let tool_call_fn = {
|
||||
// wrapper function have the same sig:
|
||||
// async fn #tool_tool_call(context: rmcp::handler::server::tool::ToolCallContext<'_, Self>)
|
||||
// -> std::result::Result<rmcp::model::CallToolResult, rmcp::Error>
|
||||
//
|
||||
// and the block part should be like:
|
||||
// {
|
||||
// use rmcp::handler::server::tool::*;
|
||||
// let (t0, context) = <T0>::from_tool_call_context_part(context)?;
|
||||
// let (t1, context) = <T1>::from_tool_call_context_part(context)?;
|
||||
// ...
|
||||
// let (tn, context) = <Tn>::from_tool_call_context_part(context)?;
|
||||
// // for params
|
||||
// ... expand helper types here
|
||||
// let (__rmcp_tool_req, context) = rmcp::model::JsonObject::from_tool_call_context_part(context)?;
|
||||
// let __#TOOL_ToolCallParam { param_0, param_1, param_2, .. } = parse_json_object(__rmcp_tool_req)?;
|
||||
// // for aggr
|
||||
// let (Parameters(aggr), context) = <Parameters<AggrType>>::from_tool_call_context_part(context)?;
|
||||
// Self::#tool_ident(to, param_0, t1, param_1, ..., param_2, tn, aggr).await.into_call_tool_result()
|
||||
//
|
||||
// }
|
||||
//
|
||||
//
|
||||
//
|
||||
|
||||
// for receiver type, name it as __rmcp_tool_receiver
|
||||
let is_async = input_fn.sig.asyncness.is_some();
|
||||
let receiver_ident = || Ident::new("__rmcp_tool_receiver", proc_macro2::Span::call_site());
|
||||
// generate the extraction part for trivial args
|
||||
let trivial_args = input_fn
|
||||
.sig
|
||||
.inputs
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(index, arg)| {
|
||||
if unextractable_args_indexes.contains(&index) {
|
||||
None
|
||||
} else {
|
||||
// get ident/type pair
|
||||
let line = match arg {
|
||||
FnArg::Typed(pat_type) => {
|
||||
let pat = &pat_type.pat;
|
||||
let ty = &pat_type.ty;
|
||||
quote! {
|
||||
let (#pat, context) = <#ty>::from_tool_call_context_part(context)?;
|
||||
}
|
||||
}
|
||||
FnArg::Receiver(r) => {
|
||||
let ty = r.ty.clone();
|
||||
let pat = receiver_ident();
|
||||
quote! {
|
||||
let (#pat, context) = <#ty>::from_tool_call_context_part(context)?;
|
||||
}
|
||||
}
|
||||
};
|
||||
Some(line)
|
||||
}
|
||||
});
|
||||
let trivial_arg_extraction_part = quote! {
|
||||
#(#trivial_args)*
|
||||
};
|
||||
let processed_arg_extraction_part = match &mut tool_macro_attrs.params {
|
||||
ToolParams::Aggregated { rust_type } => {
|
||||
let PatType { pat, ty, .. } = rust_type;
|
||||
quote! {
|
||||
let (Parameters(#pat), context) = <Parameters<#ty>>::from_tool_call_context_part(context)?;
|
||||
}
|
||||
}
|
||||
ToolParams::Params { attrs } => {
|
||||
let (param_type, temp_param_type_name) =
|
||||
create_request_type(attrs, input_fn.sig.ident.to_string());
|
||||
|
||||
let params_ident = attrs.iter().map(|attr| &attr.ident).collect::<Vec<_>>();
|
||||
quote! {
|
||||
#param_type
|
||||
let (__rmcp_tool_req, context) = rmcp::model::JsonObject::from_tool_call_context_part(context)?;
|
||||
let #temp_param_type_name {
|
||||
#(#params_ident,)*
|
||||
} = parse_json_object(__rmcp_tool_req)?;
|
||||
}
|
||||
}
|
||||
ToolParams::NoParam => {
|
||||
quote! {}
|
||||
}
|
||||
};
|
||||
// generate the execution part
|
||||
// has receiver?
|
||||
let params = &input_fn
|
||||
.sig
|
||||
.inputs
|
||||
.iter()
|
||||
.map(|fn_arg| match fn_arg {
|
||||
FnArg::Receiver(_) => {
|
||||
let pat = receiver_ident();
|
||||
quote! { #pat }
|
||||
}
|
||||
FnArg::Typed(pat_type) => {
|
||||
let pat = &pat_type.pat.clone();
|
||||
quote! { #pat }
|
||||
}
|
||||
let annotations_expr = if let Some(annotations) = attribute.annotations {
|
||||
let ToolAnnotationsAttribute {
|
||||
title,
|
||||
read_only_hint,
|
||||
destructive_hint,
|
||||
idempotent_hint,
|
||||
open_world_hint,
|
||||
} = annotations;
|
||||
fn wrap_option<T: ToTokens>(x: Option<T>) -> TokenStream {
|
||||
x.map(|x| quote! {Some(#x.into())})
|
||||
.unwrap_or(quote! { None })
|
||||
}
|
||||
let title = wrap_option(title);
|
||||
let read_only_hint = wrap_option(read_only_hint);
|
||||
let destructive_hint = wrap_option(destructive_hint);
|
||||
let idempotent_hint = wrap_option(idempotent_hint);
|
||||
let open_world_hint = wrap_option(open_world_hint);
|
||||
let token_stream = quote! {
|
||||
Some(rmcp::model::ToolAnnotations {
|
||||
title: #title,
|
||||
read_only_hint: #read_only_hint,
|
||||
destructive_hint: #destructive_hint,
|
||||
idempotent_hint: #idempotent_hint,
|
||||
open_world_hint: #open_world_hint,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let raw_fn_ident = &input_fn.sig.ident;
|
||||
let call = if is_async {
|
||||
quote! {
|
||||
Self::#raw_fn_ident(#(#params),*).await.into_call_tool_result()
|
||||
}
|
||||
} else {
|
||||
quote! {
|
||||
Self::#raw_fn_ident(#(#params),*).into_call_tool_result()
|
||||
}
|
||||
};
|
||||
// assemble the whole function
|
||||
let tool_call_fn_ident = Ident::new(
|
||||
&format!("{}_tool_call", input_fn.sig.ident),
|
||||
proc_macro2::Span::call_site(),
|
||||
);
|
||||
let raw_fn_vis = tool_macro_attrs
|
||||
.fn_item
|
||||
.vis
|
||||
.as_ref()
|
||||
.unwrap_or(&input_fn.vis);
|
||||
let raw_fn_attr = &input_fn
|
||||
.attrs
|
||||
.iter()
|
||||
.filter(|attr| !attr.path().is_ident(TOOL_IDENT))
|
||||
.collect::<Vec<_>>();
|
||||
quote! {
|
||||
#(#raw_fn_attr)*
|
||||
#raw_fn_vis async fn #tool_call_fn_ident(context: rmcp::handler::server::tool::ToolCallContext<'_, Self>)
|
||||
-> std::result::Result<rmcp::model::CallToolResult, rmcp::Error> {
|
||||
use rmcp::handler::server::tool::*;
|
||||
#trivial_arg_extraction_part
|
||||
#processed_arg_extraction_part
|
||||
#call
|
||||
}
|
||||
}
|
||||
syn::parse2::<Expr>(token_stream)?
|
||||
} else {
|
||||
none_expr()
|
||||
};
|
||||
let resolved_tool_attr = ResolvedToolAttribute {
|
||||
name: attribute.name.unwrap_or_else(|| fn_ident.to_string()),
|
||||
description: attribute
|
||||
.description
|
||||
.or_else(|| fn_item.attrs.iter().fold(None, extract_doc_line)),
|
||||
input_schema: input_schema_expr,
|
||||
annotations: annotations_expr,
|
||||
};
|
||||
let tool_attr_fn = resolved_tool_attr.into_fn(tool_attr_fn_ident)?;
|
||||
// modify the the input function
|
||||
if fn_item.sig.asyncness.is_some() {
|
||||
// 1. remove asyncness from sig
|
||||
// 2. make return type: `std::pin::Pin<Box<dyn Future<Output = #ReturnType> + Send + '_>>`
|
||||
// 3. make body: { Box::pin(async move { #body }) }
|
||||
let new_output = syn::parse2::<ReturnType>({
|
||||
match &fn_item.sig.output {
|
||||
syn::ReturnType::Default => {
|
||||
quote! { -> std::pin::Pin<Box<dyn Future<Output = ()> + Send + '_>> }
|
||||
}
|
||||
syn::ReturnType::Type(_, ty) => {
|
||||
quote! { -> std::pin::Pin<Box<dyn Future<Output = #ty> + Send + '_>> }
|
||||
}
|
||||
}
|
||||
})?;
|
||||
let prev_block = &fn_item.block;
|
||||
let new_block = syn::parse2::<syn::Block>(quote! {
|
||||
{ Box::pin(async move #prev_block ) }
|
||||
})?;
|
||||
fn_item.sig.asyncness = None;
|
||||
fn_item.sig.output = new_output;
|
||||
fn_item.block = new_block;
|
||||
}
|
||||
Ok(quote! {
|
||||
#tool_attr_fn
|
||||
#tool_call_fn
|
||||
#input_fn
|
||||
#fn_item
|
||||
})
|
||||
}
|
||||
|
||||
fn create_request_type(attrs: &[ToolFnParamAttrs], tool_name: String) -> (TokenStream, Ident) {
|
||||
let pascal_case_tool_name = tool_name.to_ascii_uppercase();
|
||||
let temp_param_type_name = Ident::new(
|
||||
&format!("__{pascal_case_tool_name}ToolCallParam",),
|
||||
proc_macro2::Span::call_site(),
|
||||
);
|
||||
(
|
||||
quote! {
|
||||
use rmcp::{serde, schemars};
|
||||
#[derive(serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
|
||||
pub struct #temp_param_type_name {
|
||||
#(#attrs)*
|
||||
}
|
||||
},
|
||||
temp_param_type_name,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
#[test]
|
||||
fn test_tool_sync_macro() -> syn::Result<()> {
|
||||
let attr = quote! {
|
||||
name = "test_tool",
|
||||
description = "test tool",
|
||||
vis =
|
||||
};
|
||||
let input = quote! {
|
||||
fn sum(&self, #[tool(aggr)] req: StructRequest) -> Result<CallToolResult, McpError> {
|
||||
Ok(CallToolResult::success(vec![Content::text((req.a + req.b).to_string())]))
|
||||
}
|
||||
};
|
||||
let input = tool(attr, input)?;
|
||||
|
||||
println!("input: {:#}", input);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_trait_tool_macro() -> syn::Result<()> {
|
||||
let attr = quote! {
|
||||
tool_box = Calculator
|
||||
name = "direct-annotated-tool",
|
||||
annotations(title = "Annotated Tool", read_only_hint = true)
|
||||
};
|
||||
let input = quote! {
|
||||
impl ServerHandler for Calculator {
|
||||
#[tool]
|
||||
fn get_info(&self) -> ServerInfo {
|
||||
ServerInfo {
|
||||
instructions: Some("A simple calculator".into()),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
async fn async_method(&self, Parameters(Request { fields }): Parameters<Request>) {
|
||||
drop(fields)
|
||||
}
|
||||
};
|
||||
let input = tool(attr, input)?;
|
||||
let _input = tool(attr, input)?;
|
||||
|
||||
println!("input: {:#}", input);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_doc_comment_description() -> syn::Result<()> {
|
||||
let attr = quote! {}; // No explicit description
|
||||
|
|
@ -764,6 +268,7 @@ mod test {
|
|||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_explicit_description_priority() -> syn::Result<()> {
|
||||
let attr = quote! {
|
||||
|
|
|
|||
51
crates/rmcp-macros/src/tool_handler.rs
Normal file
51
crates/rmcp-macros/src/tool_handler.rs
Normal file
|
|
@ -0,0 +1,51 @@
|
|||
use darling::{FromMeta, ast::NestedMeta};
|
||||
use proc_macro2::TokenStream;
|
||||
use quote::{ToTokens, quote};
|
||||
use syn::{Expr, ImplItem, ItemImpl};
|
||||
|
||||
#[derive(FromMeta)]
|
||||
#[darling(default)]
|
||||
pub struct ToolHandlerAttribute {
|
||||
pub router: Expr,
|
||||
}
|
||||
|
||||
impl Default for ToolHandlerAttribute {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
router: syn::parse2(quote! {
|
||||
self.tool_router
|
||||
})
|
||||
.unwrap(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn tool_handler(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
|
||||
let attr_args = NestedMeta::parse_meta_list(attr)?;
|
||||
let ToolHandlerAttribute { router } = ToolHandlerAttribute::from_list(&attr_args)?;
|
||||
let mut item_impl = syn::parse2::<ItemImpl>(input.clone())?;
|
||||
let tool_call_fn = quote! {
|
||||
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);
|
||||
#router.call(tcc).await
|
||||
}
|
||||
};
|
||||
let tool_list_fn = quote! {
|
||||
async fn list_tools(
|
||||
&self,
|
||||
_request: Option<rmcp::model::PaginatedRequestParam>,
|
||||
_context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
||||
) -> Result<rmcp::model::ListToolsResult, rmcp::Error> {
|
||||
Ok(rmcp::model::ListToolsResult::with_all_items(#router.list_all()))
|
||||
}
|
||||
};
|
||||
let tool_call_fn = syn::parse2::<ImplItem>(tool_call_fn)?;
|
||||
let tool_list_fn = syn::parse2::<ImplItem>(tool_list_fn)?;
|
||||
item_impl.items.push(tool_call_fn);
|
||||
item_impl.items.push(tool_list_fn);
|
||||
Ok(item_impl.into_token_stream())
|
||||
}
|
||||
90
crates/rmcp-macros/src/tool_router.rs
Normal file
90
crates/rmcp-macros/src/tool_router.rs
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
//! ```ignore
|
||||
//! #[rmcp::tool_router(router)]
|
||||
//! impl Handler {
|
||||
//!
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
|
||||
use darling::{FromMeta, ast::NestedMeta};
|
||||
use proc_macro2::TokenStream;
|
||||
use quote::{ToTokens, format_ident, quote};
|
||||
use syn::{Ident, ImplItem, ItemImpl, Visibility};
|
||||
|
||||
#[derive(FromMeta)]
|
||||
#[darling(default)]
|
||||
pub struct ToolRouterAttribute {
|
||||
pub router: Ident,
|
||||
pub vis: Option<Visibility>,
|
||||
}
|
||||
|
||||
impl Default for ToolRouterAttribute {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
router: format_ident!("tool_router"),
|
||||
vis: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn tool_router(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
|
||||
let attr_args = NestedMeta::parse_meta_list(attr)?;
|
||||
let ToolRouterAttribute { router, vis } = ToolRouterAttribute::from_list(&attr_args)?;
|
||||
let mut item_impl = syn::parse2::<ItemImpl>(input.clone())?;
|
||||
// find all function marked with `#[rmcp::tool]`
|
||||
let tool_attr_fns: Vec<_> = item_impl
|
||||
.items
|
||||
.iter()
|
||||
.filter_map(|item| {
|
||||
if let syn::ImplItem::Fn(fn_item) = item {
|
||||
fn_item
|
||||
.attrs
|
||||
.iter()
|
||||
.any(|attr| {
|
||||
attr.path()
|
||||
.segments
|
||||
.last()
|
||||
.is_some_and(|seg| seg.ident == "tool")
|
||||
})
|
||||
.then_some(&fn_item.sig.ident)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let mut routers = vec![];
|
||||
for handler in tool_attr_fns {
|
||||
let tool_attr_fn_ident = format_ident!("{handler}_tool_attr");
|
||||
routers.push(quote! {
|
||||
.with_route((Self::#tool_attr_fn_ident(), Self::#handler))
|
||||
})
|
||||
}
|
||||
let router_fn = syn::parse2::<ImplItem>(quote! {
|
||||
#vis fn #router() -> rmcp::handler::server::router::tool::ToolRouter<Self> {
|
||||
rmcp::handler::server::router::tool::ToolRouter::<Self>::new()
|
||||
#(#routers)*
|
||||
}
|
||||
})?;
|
||||
item_impl.items.push(router_fn);
|
||||
Ok(item_impl.into_token_stream())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
#[test]
|
||||
fn test_router_attr() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let attr = quote! {
|
||||
router = test_router,
|
||||
};
|
||||
let attr_args = NestedMeta::parse_meta_list(attr)?;
|
||||
let ToolRouterAttribute { router, vis } = ToolRouterAttribute::from_list(&attr_args)?;
|
||||
println!("router: {}", router);
|
||||
if let Some(vis) = vis {
|
||||
println!("visibility: {}", vis.to_token_stream());
|
||||
} else {
|
||||
println!("visibility: None");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue