fastmcp_rs_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::{format_ident, quote};
3use syn::{parse_macro_input, ItemFn, Meta, Lit, Expr, MetaNameValue};
4use syn::parse::Parser;
5
6fn parse_args(attr: proc_macro2::TokenStream) -> Vec<Meta> {
7    let parser = syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated;
8    parser.parse2(attr).map(|p| p.into_iter().collect()).unwrap_or_default()
9}
10
11fn parse_kv_string(metas: &[Meta], key: &str) -> Option<String> {
12    for m in metas {
13        if let Meta::NameValue(nv) = m {
14            if nv.path.is_ident(key) {
15                if let Expr::Lit(expr_lit) = &nv.value {
16                    if let Lit::Str(s) = &expr_lit.lit {
17                        return Some(s.value());
18                    }
19                }
20            }
21        }
22    }
23    None
24}
25
26fn parse_kv_expr(metas: &[Meta], key: &str) -> Option<Expr> {
27    for m in metas {
28        if let Meta::NameValue(nv) = m {
29            if nv.path.is_ident(key) {
30                return Some(nv.value.clone());
31            }
32        }
33    }
34    None
35}
36
37fn parse_parameters_json(metas: &[Meta]) -> Option<proc_macro2::TokenStream> {
38    for m in metas {
39        if let Meta::List(list) = m {
40            if list.path.is_ident("parameters_json") {
41                return Some(list.tokens.clone());
42            }
43        }
44    }
45    None
46}
47
48fn parse_annotations(metas: &[Meta]) -> Option<Vec<(String, Expr)>> {
49    for m in metas {
50        if let Meta::List(list) = m {
51            if list.path.is_ident("annotations") {
52                let punct = syn::punctuated::Punctuated::<MetaNameValue, syn::Token![,]>::parse_terminated
53                    .parse2(list.tokens.clone())
54                    .unwrap_or_default();
55                let mut pairs = Vec::new();
56                for nv in punct {
57                    if let Some(ident) = nv.path.get_ident() {
58                        let key = ident.to_string();
59                        let val = nv.value.clone();
60                        pairs.push((key, val));
61                    }
62                }
63                return Some(pairs);
64            }
65        }
66    }
67    None
68}
69
70#[proc_macro_attribute]
71pub fn mcp_tool(attr: TokenStream, item: TokenStream) -> TokenStream {
72    let metas = parse_args(attr.into());
73    let input_fn = parse_macro_input!(item as ItemFn);
74
75    let fn_name = input_fn.sig.ident.clone();
76    let tool_fn_name = format_ident!("{}_tool", fn_name);
77    let register_fn_name = format_ident!("{}_register", fn_name);
78    let factory_static = format_ident!("__{}_factory", fn_name);
79
80    // Required: name
81    let name = parse_kv_string(&metas, "name").unwrap_or_else(|| fn_name.to_string());
82
83    let description = parse_kv_string(&metas, "description");
84    let summary = parse_kv_string(&metas, "summary");
85    let parameters_expr = parse_kv_expr(&metas, "parameters");
86    let parameters_json = parse_parameters_json(&metas);
87    let annotations = parse_annotations(&metas);
88
89    // Prepare optional sections
90    let description_set = description.as_ref().map(|s| quote! { __def = __def.with_description(#s); });
91    let summary_set = summary.as_ref().map(|s| quote! { __def = __def.with_summary(#s); });
92    let parameters_set = parameters_expr.as_ref().map(|e| quote! { __def = __def.with_parameters(#e); });
93    let parameters_json_set = parameters_json.as_ref().map(|ts| quote! { __def = __def.with_parameters(serde_json::json!({ #ts })); });
94    let annotations_set = annotations.as_ref().map(|pairs| {
95        let inserts: Vec<proc_macro2::TokenStream> = pairs.iter()
96            .map(|(k, v)| quote! { __ann.insert(String::from(#k), #v); })
97            .collect();
98        quote! {
99            let mut __ann = serde_json::Map::<String, serde_json::Value>::new();
100            #( #inserts )*
101            __def = __def.with_annotations(__ann);
102        }
103    });
104
105    let output_tokens = quote! {
106        #input_fn
107
108        #[allow(non_snake_case)]
109        pub fn #tool_fn_name() -> fastmcp_rs::tool::ToolDefinition {
110            let mut __def = fastmcp_rs::tool::ToolDefinition::new(#name, #fn_name);
111            #description_set
112            #summary_set
113            #parameters_set
114            #parameters_json_set
115            #annotations_set
116            __def
117        }
118
119        #[allow(non_snake_case)]
120        pub fn #register_fn_name(server: &fastmcp_rs::FastMcpServer) -> fastmcp_rs::Result<()> {
121            server.register_tool(#tool_fn_name())
122        }
123
124        #[cfg(feature = "auto-register")]
125        #[linkme::distributed_slice(fastmcp_rs::tool::MCP_TOOL_FACTORIES)]
126        pub static #factory_static: fastmcp_rs::tool::ToolFactory = || { #tool_fn_name() };
127    };
128
129    output_tokens.into()
130}