mcp_attr_macros/
lib.rs

1#![allow(unused)]
2
3use std::{
4    collections::{HashMap, HashSet},
5    str::FromStr,
6};
7
8use proc_macro2::TokenStream;
9use quote::{ToTokens, quote};
10use structmeta::{NameArgs, NameValue, StructMeta};
11use syn::{
12    Attribute, Error, FnArg, Ident, ImplItem, ImplItemFn, ItemFn, ItemImpl, LitStr, Pat, Path,
13    Result, Type, parse::Parse, parse2, spanned::Spanned,
14};
15use uri_template_ex::UriTemplate;
16
17use syn_utils::{get_element, is_path, is_type};
18use utils::{get_trait_path, is_defined};
19
20use crate::prompts::{PromptAttr, PromptEntry};
21use crate::resources::{ResourceAttr, ResourceEntry};
22use crate::tools::{ToolAttr, ToolEntry};
23use crate::utils::{build_if, drain_attr};
24
25#[macro_use]
26mod syn_utils;
27mod utils;
28
29mod prompts;
30mod resources;
31mod tools;
32
33#[proc_macro_attribute]
34pub fn mcp_server(
35    attr: proc_macro::TokenStream,
36    item: proc_macro::TokenStream,
37) -> proc_macro::TokenStream {
38    let mut item: TokenStream = item.into();
39    let mut es = Vec::new();
40    match build(attr.into(), item.clone(), &mut es) {
41        Ok(mut s) => {
42            for e in es {
43                s.extend(e.to_compile_error());
44            }
45            s
46        }
47        Err(e) => e.to_compile_error(),
48    }
49    .into()
50}
51
52fn build(attr: TokenStream, item: TokenStream, es: &mut Vec<Error>) -> Result<TokenStream> {
53    let mut item_impl: ItemImpl = parse2(item)?;
54    let mut attr: McpAttr = parse2(attr)?;
55    let trait_path = get_trait_path(&item_impl)?.clone();
56    if item_impl.unsafety.is_some() {
57        bail!(item_impl.span(), "Unsafe is not allowed");
58    }
59    if item_impl.defaultness.is_some() {
60        bail!(item_impl.span(), "Default is not allowed");
61    }
62    let is_defined_resources_list = is_defined(&item_impl.items, "resources_list");
63    let mut b = McpBuilder::new();
64    let mut items_trait = Vec::new();
65    let mut items_type = Vec::new();
66    for mut item in item_impl.items {
67        match b.push(&mut item) {
68            Ok(true) => items_type.push(item),
69            Ok(false) => items_trait.push(item),
70            Err(e) => {
71                items_type.push(item);
72                es.push(e);
73            }
74        }
75    }
76    let b = b.build(&items_trait)?;
77    let (impl_generics, ty_generics, where_clause) = item_impl.generics.split_for_impl();
78
79    let self_ty = &item_impl.self_ty;
80    let attrs = &item_impl.attrs;
81    let ts = quote! {
82        #[automatically_derived]
83        #(#attrs)*
84        impl<#impl_generics> #trait_path for #self_ty #ty_generics #where_clause {
85            #(#items_trait)*
86            #b
87        }
88
89        #[automatically_derived]
90        #(#attrs)*
91        impl<#impl_generics> #self_ty #ty_generics #where_clause {
92            #(#items_type)*
93        }
94    };
95    if attr.dump {
96        panic!("// ===== start generated code =====\n{ts}\n// ===== end generated code =====\n");
97    }
98    Ok(ts)
99}
100
101struct McpBuilder {
102    prompts: Vec<PromptEntry>,
103    resources: Vec<ResourceEntry>,
104    tools: Vec<ToolEntry>,
105}
106
107impl McpBuilder {
108    fn new() -> Self {
109        Self {
110            prompts: Vec::new(),
111            resources: Vec::new(),
112            tools: Vec::new(),
113        }
114    }
115    fn push(&mut self, item: &mut ImplItem) -> Result<bool> {
116        if let ImplItem::Fn(f) = item {
117            let Some(attr) = drain_attr(&mut f.attrs)? else {
118                return Ok(false);
119            };
120            match attr {
121                ItemAttr::Prompt(attr) => self.prompts.push(PromptEntry::new(f, attr)?),
122                ItemAttr::Resource(attr) => self.resources.push(ResourceEntry::new(f, attr)?),
123                ItemAttr::Tool(attr) => self.tools.push(ToolEntry::new(f, attr)?),
124            }
125            return Ok(true);
126        }
127        Ok(false)
128    }
129
130    fn build(&self, items: &[ImplItem]) -> Result<TokenStream> {
131        let capabilities = build_if(!is_defined(items, "capabilities"), || {
132            self.build_capabilities(items)
133        })?;
134        let prompts = build_if(!self.prompts.is_empty(), || self.build_prompts())?;
135        let resources = build_if(!self.resources.is_empty(), || self.build_resources(items))?;
136        let tools = build_if(!self.tools.is_empty(), || self.build_tools())?;
137        Ok(quote! {
138            #capabilities
139            #prompts
140            #resources
141            #tools
142        })
143    }
144    fn build_capabilities(&self, items: &[ImplItem]) -> Result<TokenStream> {
145        let prompts = if !self.prompts.is_empty() || is_defined(items, "prompts_list") {
146            quote!(Some(::mcp_attr::schema::ServerCapabilitiesPrompts {
147                ..::std::default::Default::default()
148            }))
149        } else {
150            quote!(None)
151        };
152        let resources = if !self.resources.is_empty() || is_defined(items, "resources_read") {
153            quote!(Some(::mcp_attr::schema::ServerCapabilitiesResources {
154                ..::std::default::Default::default()
155            }))
156        } else {
157            quote!(None)
158        };
159        let tools = if !self.tools.is_empty() || is_defined(items, "tools_list") {
160            quote!(Some(::mcp_attr::schema::ServerCapabilitiesTools {
161                ..::std::default::Default::default()
162            }))
163        } else {
164            quote!(None)
165        };
166        Ok(quote! {
167            fn capabilities(&self) -> ::mcp_attr::schema::ServerCapabilities {
168                ::mcp_attr::schema::ServerCapabilities {
169                    prompts: #prompts,
170                    resources: #resources,
171                    tools: #tools,
172                    ..::std::default::Default::default()
173                }
174            }
175        })
176    }
177    fn build_prompts(&self) -> Result<TokenStream> {
178        let list = self.build_prompts_list()?;
179        let get = self.build_prompts_get()?;
180        Ok(quote! {
181            #list
182            #get
183        })
184    }
185    fn build_resources(&self, items: &[ImplItem]) -> Result<TokenStream> {
186        let list = build_if(!is_defined(items, "resources_list"), || {
187            self.build_resources_list()
188        })?;
189        let templates_list = self.build_resources_templates_list()?;
190        let read = self.build_resources_read()?;
191        Ok(quote! {
192            #list
193            #templates_list
194            #read
195        })
196    }
197    fn build_tools(&self) -> Result<TokenStream> {
198        let list = self.build_tools_list()?;
199        let call = self.build_tools_call()?;
200        Ok(quote! {
201            #list
202            #call
203        })
204    }
205    fn build_prompts_list(&self) -> Result<TokenStream> {
206        PromptEntry::build_list(&self.prompts)
207    }
208    fn build_prompts_get(&self) -> Result<TokenStream> {
209        PromptEntry::build_get(&self.prompts)
210    }
211    fn build_resources_list(&self) -> Result<TokenStream> {
212        ResourceEntry::build_list(&self.resources)
213    }
214    fn build_resources_templates_list(&self) -> Result<TokenStream> {
215        ResourceEntry::build_templates_list(&self.resources)
216    }
217    fn build_resources_read(&self) -> Result<TokenStream> {
218        ResourceEntry::build_read(&self.resources)
219    }
220
221    fn build_tools_list(&self) -> Result<TokenStream> {
222        ToolEntry::build_list(&self.tools)
223    }
224    fn build_tools_call(&self) -> Result<TokenStream> {
225        ToolEntry::build_call(&self.tools)
226    }
227}
228
229#[derive(StructMeta, Default)]
230struct McpAttr {
231    dump: bool,
232}
233
234enum ItemAttr {
235    Prompt(PromptAttr),
236    Resource(ResourceAttr),
237    Tool(ToolAttr),
238}
mcp_attr_macros/lib.rs

mcp_attr_macros/
lib.rs
