mcp_attr_macros/
lib.rs

1#![allow(unused)]
2
3use std::{
4    collections::{HashMap, HashSet},
5    str::FromStr,
6};
7
8use proc_macro2::{Span, TokenStream};
9use quote::{ToTokens, format_ident, quote};
10use structmeta::{NameArgs, NameValue, StructMeta};
11use syn::{
12    Attribute, Error, FnArg, Ident, ImplItem, ImplItemFn, ItemFn, ItemImpl, LitStr, Pat, Path,
13    Result, Token, Type,
14    parse::{Parse, ParseStream},
15    parse2,
16    punctuated::Punctuated,
17    spanned::Spanned,
18};
19use uri_template_ex::UriTemplate;
20
21use syn_utils::{get_element, into_macro_output, is_path, is_type};
22use utils::{get_trait_path, is_defined};
23
24use crate::prompts::{PromptAttr, PromptEntry};
25use crate::resources::{ResourceAttr, ResourceEntry};
26use crate::tools::{ToolAttr, ToolEntry};
27use crate::utils::{build_if, drain_attr};
28
29#[macro_use]
30mod syn_utils;
31mod utils;
32
33mod prompts;
34mod resources;
35mod tools;
36
37#[proc_macro_attribute]
38pub fn mcp_server(
39    attr: proc_macro::TokenStream,
40    item: proc_macro::TokenStream,
41) -> proc_macro::TokenStream {
42    let mut item: TokenStream = item.into();
43    let mut es = Vec::new();
44    match build_mcp_server(attr.into(), item.clone(), &mut es) {
45        Ok(mut s) => {
46            for e in es {
47                s.extend(e.to_compile_error());
48            }
49            s
50        }
51        Err(e) => e.to_compile_error(),
52    }
53    .into()
54}
55
56#[proc_macro]
57pub fn route(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
58    into_macro_output(build_route(item.into()))
59}
60
61#[proc_macro_attribute]
62pub fn tool(
63    attr: proc_macro::TokenStream,
64    item: proc_macro::TokenStream,
65) -> proc_macro::TokenStream {
66    into_macro_output(build_tool(attr.into(), item.into()))
67}
68
69#[proc_macro_attribute]
70pub fn resource(
71    attr: proc_macro::TokenStream,
72    item: proc_macro::TokenStream,
73) -> proc_macro::TokenStream {
74    into_macro_output(build_resource(attr.into(), item.into()))
75}
76
77#[proc_macro_attribute]
78pub fn prompt(
79    attr: proc_macro::TokenStream,
80    item: proc_macro::TokenStream,
81) -> proc_macro::TokenStream {
82    into_macro_output(build_prompt(attr.into(), item.into()))
83}
84
85fn build_mcp_server(
86    attr: TokenStream,
87    item: TokenStream,
88    es: &mut Vec<Error>,
89) -> Result<TokenStream> {
90    let mut item_impl: ItemImpl = parse2(item)?;
91    let mut attr: McpAttr = parse2(attr)?;
92    let trait_path = get_trait_path(&item_impl)?.clone();
93    if item_impl.unsafety.is_some() {
94        bail!(item_impl.span(), "Unsafe is not allowed");
95    }
96    if item_impl.defaultness.is_some() {
97        bail!(item_impl.span(), "Default is not allowed");
98    }
99    let is_defined_resources_list = is_defined(&item_impl.items, "resources_list");
100    let mut b = McpBuilder::new();
101    let mut items_trait = Vec::new();
102    let mut items_type = Vec::new();
103    for mut item in item_impl.items {
104        match b.push(&mut item) {
105            Ok(true) => items_type.push(item),
106            Ok(false) => items_trait.push(item),
107            Err(e) => {
108                items_type.push(item);
109                es.push(e);
110            }
111        }
112    }
113    let b = b.build(&items_trait)?;
114    let (impl_generics, ty_generics, where_clause) = item_impl.generics.split_for_impl();
115
116    let self_ty = &item_impl.self_ty;
117    let attrs = &item_impl.attrs;
118    let ts = quote! {
119        #[automatically_derived]
120        #(#attrs)*
121        impl<#impl_generics> #trait_path for #self_ty #ty_generics #where_clause {
122            #(#items_trait)*
123            #b
124        }
125
126        #[automatically_derived]
127        #(#attrs)*
128        impl<#impl_generics> #self_ty #ty_generics #where_clause {
129            #(#items_type)*
130        }
131    };
132    if attr.dump {
133        dump_code(ts);
134    }
135    Ok(ts)
136}
137
138struct McpBuilder {
139    prompts: Vec<PromptEntry>,
140    resources: Vec<ResourceEntry>,
141    tools: Vec<ToolEntry>,
142}
143
144impl McpBuilder {
145    fn new() -> Self {
146        Self {
147            prompts: Vec::new(),
148            resources: Vec::new(),
149            tools: Vec::new(),
150        }
151    }
152    fn push(&mut self, item: &mut ImplItem) -> Result<bool> {
153        if let ImplItem::Fn(f) = item {
154            let Some(attr) = drain_attr(&mut f.attrs)? else {
155                return Ok(false);
156            };
157            match attr {
158                ItemAttr::Prompt(attr) => {
159                    self.prompts.push(PromptEntry::from_impl_item_fn(f, attr)?)
160                }
161                ItemAttr::Resource(attr) => self
162                    .resources
163                    .push(ResourceEntry::from_impl_item_fn(f, attr)?),
164                ItemAttr::Tool(attr) => self.tools.push(ToolEntry::from_impl_item_fn(f, attr)?),
165            }
166            return Ok(true);
167        }
168        Ok(false)
169    }
170
171    fn build(&self, items: &[ImplItem]) -> Result<TokenStream> {
172        let capabilities = build_if(!is_defined(items, "capabilities"), || {
173            self.build_capabilities(items)
174        })?;
175        let prompts = build_if(!self.prompts.is_empty(), || self.build_prompts())?;
176        let resources = build_if(!self.resources.is_empty(), || self.build_resources(items))?;
177        let tools = build_if(!self.tools.is_empty(), || self.build_tools())?;
178        Ok(quote! {
179            #capabilities
180            #prompts
181            #resources
182            #tools
183        })
184    }
185    fn build_capabilities(&self, items: &[ImplItem]) -> Result<TokenStream> {
186        let prompts = if !self.prompts.is_empty() || is_defined(items, "prompts_list") {
187            quote!(Some(::mcp_attr::schema::ServerCapabilitiesPrompts {
188                ..::std::default::Default::default()
189            }))
190        } else {
191            quote!(None)
192        };
193        let resources = if !self.resources.is_empty() || is_defined(items, "resources_read") {
194            quote!(Some(::mcp_attr::schema::ServerCapabilitiesResources {
195                ..::std::default::Default::default()
196            }))
197        } else {
198            quote!(None)
199        };
200        let tools = if !self.tools.is_empty() || is_defined(items, "tools_list") {
201            quote!(Some(::mcp_attr::schema::ServerCapabilitiesTools {
202                ..::std::default::Default::default()
203            }))
204        } else {
205            quote!(None)
206        };
207        Ok(quote! {
208            fn capabilities(&self) -> ::mcp_attr::schema::ServerCapabilities {
209                ::mcp_attr::schema::ServerCapabilities {
210                    prompts: #prompts,
211                    resources: #resources,
212                    tools: #tools,
213                    ..::std::default::Default::default()
214                }
215            }
216        })
217    }
218    fn build_prompts(&self) -> Result<TokenStream> {
219        let list = self.build_prompts_list()?;
220        let get = self.build_prompts_get()?;
221        Ok(quote! {
222            #list
223            #get
224        })
225    }
226    fn build_resources(&self, items: &[ImplItem]) -> Result<TokenStream> {
227        let list = build_if(!is_defined(items, "resources_list"), || {
228            self.build_resources_list()
229        })?;
230        let templates_list = self.build_resources_templates_list()?;
231        let read = self.build_resources_read()?;
232        Ok(quote! {
233            #list
234            #templates_list
235            #read
236        })
237    }
238    fn build_tools(&self) -> Result<TokenStream> {
239        let list = self.build_tools_list()?;
240        let call = self.build_tools_call()?;
241        Ok(quote! {
242            #list
243            #call
244        })
245    }
246    fn build_prompts_list(&self) -> Result<TokenStream> {
247        PromptEntry::build_list(&self.prompts)
248    }
249    fn build_prompts_get(&self) -> Result<TokenStream> {
250        PromptEntry::build_get(&self.prompts)
251    }
252    fn build_resources_list(&self) -> Result<TokenStream> {
253        ResourceEntry::build_list(&self.resources)
254    }
255    fn build_resources_templates_list(&self) -> Result<TokenStream> {
256        ResourceEntry::build_templates_list(&self.resources)
257    }
258    fn build_resources_read(&self) -> Result<TokenStream> {
259        ResourceEntry::build_read(&self.resources)
260    }
261
262    fn build_tools_list(&self) -> Result<TokenStream> {
263        ToolEntry::build_list(&self.tools)
264    }
265    fn build_tools_call(&self) -> Result<TokenStream> {
266        ToolEntry::build_call(&self.tools)
267    }
268}
269
270#[derive(StructMeta, Default)]
271struct McpAttr {
272    dump: bool,
273}
274
275enum ItemAttr {
276    Prompt(PromptAttr),
277    Resource(ResourceAttr),
278    Tool(ToolAttr),
279}
280
281fn build_route(item: TokenStream) -> Result<TokenStream> {
282    struct PathList(Punctuated<Path, Token![,]>);
283    impl Parse for PathList {
284        fn parse(input: ParseStream) -> Result<Self> {
285            Ok(Self(Punctuated::parse_terminated(input)?))
286        }
287    }
288    let mut path_list: PathList = parse2(item)?;
289    let mut exprs = Vec::new();
290    for mut path in path_list.0 {
291        let last = path.segments.last_mut().unwrap();
292        let fn_ident = last.ident.clone();
293        last.ident = route_ident(&fn_ident);
294        last.ident.set_span(Span::call_site());
295        exprs.push(quote! {
296            {
297                let _dummy = #fn_ident; // Ensure rust-analyzer can rename the function.
298                ::std::convert::Into::<::mcp_attr::server::builder::Route>::into(#path()?)
299            }
300        });
301    }
302    Ok(quote! {
303        [#(#exprs),*]
304    })
305}
306fn route_ident(ident: &Ident) -> Ident {
307    format_ident!("__route_of_{}", ident)
308}
309fn build_tool(attr: TokenStream, item: TokenStream) -> Result<TokenStream> {
310    let mut f: ItemFn = parse2(item)?;
311    let attr: ToolAttr = parse2(attr)?;
312    let dump = attr.dump;
313    let e = ToolEntry::from_item_fn(&mut f, attr)?;
314    let ret = e.build_route();
315    Ok(make_extend(f, ret, dump))
316}
317
318fn build_resource(attr: TokenStream, item: TokenStream) -> Result<TokenStream> {
319    let mut f: ItemFn = parse2(item)?;
320    let attr: ResourceAttr = parse2(attr)?;
321    let dump = attr.dump;
322    let e = ResourceEntry::from_item_fn(&mut f, attr)?;
323    let ret = e.build_route();
324    Ok(make_extend(f, ret, dump))
325}
326
327fn build_prompt(attr: TokenStream, item: TokenStream) -> Result<TokenStream> {
328    let mut f: ItemFn = parse2(item)?;
329    let attr: PromptAttr = parse2(attr)?;
330    let dump = attr.dump;
331    let e = PromptEntry::from_item_fn(&mut f, attr)?;
332    let ret = e.build_route();
333    Ok(make_extend(f, ret, dump))
334}
335
336fn make_extend(source: impl ToTokens, code: Result<TokenStream>, dump: bool) -> TokenStream {
337    let code = match code {
338        Ok(code) => {
339            if dump {
340                dump_code(code);
341            }
342            code
343        }
344        Err(e) => e.to_compile_error(),
345    };
346    quote! {
347        #source
348        #code
349    }
350}
351fn dump_code(code: TokenStream) -> ! {
352    panic!("// ===== start generated code =====\n{code}\n// ===== end generated code =====\n");
353}
mcp_attr_macros/lib.rs

mcp_attr_macros/
lib.rs
