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 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 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}