wasmcp_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, ItemMod, ItemFn, Attribute, FnArg, Type, Pat};
4use proc_macro2::TokenStream as TokenStream2;
5
6/// Mark a module as containing MCP tools.
7/// Functions in this module will be automatically registered as tools.
8#[proc_macro_attribute]
9pub fn mcp_tools(_attr: TokenStream, input: TokenStream) -> TokenStream {
10    let module = parse_macro_input!(input as ItemMod);
11    
12    // Collect all functions in the module as tools
13    let mut tool_defs = Vec::<TokenStream2>::new();
14    let mut tool_dispatches = Vec::<TokenStream2>::new();
15    
16    if let Some((_, items)) = &module.content {
17        for item in items.iter() {
18            if let syn::Item::Fn(func) = item {
19                let (def, dispatch) = generate_tool_info(func);
20                tool_defs.push(def);
21                tool_dispatches.push(dispatch);
22            }
23        }
24    }
25    
26    let module_name = &module.ident;
27    
28    // Generate the component and tool handler implementation
29    quote! {
30        #module
31        
32        pub struct Component;
33        
34        // Core implementation (always required)
35        impl crate::bindings::exports::fastertools::mcp::core::Guest for Component {
36            fn handle_initialize(_request: crate::bindings::fastertools::mcp::session::InitializeRequest) 
37                -> Result<crate::bindings::fastertools::mcp::session::InitializeResponse, crate::bindings::fastertools::mcp::types::McpError> {
38                Ok(crate::bindings::fastertools::mcp::session::InitializeResponse {
39                    protocol_version: "0.1.0".to_string(),
40                    capabilities: crate::bindings::fastertools::mcp::session::ServerCapabilities {
41                        tools: Some(crate::bindings::fastertools::mcp::session::ToolsCapability { 
42                            list_changed: Some(false) 
43                        }),
44                        resources: None,
45                        prompts: None,
46                        experimental: None,
47                        logging: None,
48                        completions: None,
49                    },
50                    server_info: crate::bindings::fastertools::mcp::session::ImplementationInfo {
51                        name: env!("CARGO_PKG_NAME").to_string(),
52                        version: env!("CARGO_PKG_VERSION").to_string(),
53                        title: Some(env!("CARGO_PKG_NAME").to_string()),
54                    },
55                    instructions: None,
56                    meta: None,
57                })
58            }
59            
60            fn handle_initialized() -> Result<(), crate::bindings::fastertools::mcp::types::McpError> {
61                Ok(())
62            }
63            
64            fn handle_ping() -> Result<(), crate::bindings::fastertools::mcp::types::McpError> {
65                Ok(())
66            }
67            
68            fn handle_shutdown() -> Result<(), crate::bindings::fastertools::mcp::types::McpError> {
69                Ok(())
70            }
71        }
72        
73        // Tool handler implementation
74        impl crate::bindings::exports::fastertools::mcp::tool_handler::Guest for Component {
75            fn handle_list_tools(_request: crate::bindings::fastertools::mcp::tools::ListToolsRequest) 
76                -> Result<crate::bindings::fastertools::mcp::tools::ListToolsResponse, crate::bindings::fastertools::mcp::types::McpError> {
77                Ok(crate::bindings::fastertools::mcp::tools::ListToolsResponse {
78                    tools: vec![#(#tool_defs),*],
79                    next_cursor: None,
80                    meta: None,
81                })
82            }
83            
84            fn handle_call_tool(request: crate::bindings::fastertools::mcp::tools::CallToolRequest) 
85                -> Result<crate::bindings::fastertools::mcp::tools::ToolResult, crate::bindings::fastertools::mcp::types::McpError> {
86                let args = if let Some(args_str) = &request.arguments {
87                    serde_json::from_str(args_str)
88                        .map_err(|e| crate::bindings::fastertools::mcp::types::McpError {
89                            code: crate::bindings::fastertools::mcp::types::ErrorCode::InvalidParams,
90                            message: format!("Invalid arguments: {}", e),
91                            data: None,
92                        })?
93                } else {
94                    serde_json::Value::Object(serde_json::Map::new())
95                };
96                
97                use #module_name::*;
98                
99                match request.name.as_str() {
100                    #(#tool_dispatches)*
101                    _ => Err(crate::bindings::fastertools::mcp::types::McpError {
102                        code: crate::bindings::fastertools::mcp::types::ErrorCode::ToolNotFound,
103                        message: format!("Unknown tool: {}", request.name),
104                        data: None,
105                    })
106                }
107            }
108        }
109        
110        // Empty stub implementations for unused capabilities
111        impl crate::bindings::exports::fastertools::mcp::resource_handler::Guest for Component {
112            fn handle_list_resources(_request: crate::bindings::fastertools::mcp::resources::ListResourcesRequest) 
113                -> Result<crate::bindings::fastertools::mcp::resources::ListResourcesResponse, crate::bindings::fastertools::mcp::types::McpError> {
114                Ok(crate::bindings::fastertools::mcp::resources::ListResourcesResponse {
115                    resources: vec![],
116                    next_cursor: None,
117                    meta: None,
118                })
119            }
120            
121            fn handle_list_resource_templates(_request: crate::bindings::fastertools::mcp::resources::ListTemplatesRequest) 
122                -> Result<crate::bindings::fastertools::mcp::resources::ListTemplatesResponse, crate::bindings::fastertools::mcp::types::McpError> {
123                Ok(crate::bindings::fastertools::mcp::resources::ListTemplatesResponse {
124                    templates: vec![],
125                    next_cursor: None,
126                    meta: None,
127                })
128            }
129            
130            fn handle_read_resource(_request: crate::bindings::fastertools::mcp::resources::ReadResourceRequest) 
131                -> Result<crate::bindings::fastertools::mcp::resources::ReadResourceResponse, crate::bindings::fastertools::mcp::types::McpError> {
132                Err(crate::bindings::fastertools::mcp::types::McpError {
133                    code: crate::bindings::fastertools::mcp::types::ErrorCode::ResourceNotFound,
134                    message: "This server does not provide resources".to_string(),
135                    data: None,
136                })
137            }
138            
139            fn handle_subscribe_resource(_request: crate::bindings::fastertools::mcp::resources::SubscribeRequest) 
140                -> Result<(), crate::bindings::fastertools::mcp::types::McpError> {
141                Ok(())
142            }
143            
144            fn handle_unsubscribe_resource(_request: crate::bindings::fastertools::mcp::resources::UnsubscribeRequest) 
145                -> Result<(), crate::bindings::fastertools::mcp::types::McpError> {
146                Ok(())
147            }
148        }
149        
150        // Prompt handler implementation (stub)
151        impl crate::bindings::exports::fastertools::mcp::prompt_handler::Guest for Component {
152            fn handle_list_prompts(_request: crate::bindings::fastertools::mcp::prompts::ListPromptsRequest) 
153                -> Result<crate::bindings::fastertools::mcp::prompts::ListPromptsResponse, crate::bindings::fastertools::mcp::types::McpError> {
154                Ok(crate::bindings::fastertools::mcp::prompts::ListPromptsResponse {
155                    prompts: vec![],
156                    next_cursor: None,
157                    meta: None,
158                })
159            }
160            
161            fn handle_get_prompt(_request: crate::bindings::fastertools::mcp::prompts::GetPromptRequest) 
162                -> Result<crate::bindings::fastertools::mcp::prompts::GetPromptResponse, crate::bindings::fastertools::mcp::types::McpError> {
163                Err(crate::bindings::fastertools::mcp::types::McpError {
164                    code: crate::bindings::fastertools::mcp::types::ErrorCode::PromptNotFound,
165                    message: "This server does not provide prompts".to_string(),
166                    data: None,
167                })
168            }
169        }
170    }.into()
171}
172
173/// Mark a module as containing MCP resources.
174/// Functions in this module will be automatically registered as resources.
175#[proc_macro_attribute]
176pub fn mcp_resources(_attr: TokenStream, input: TokenStream) -> TokenStream {
177    let module = parse_macro_input!(input as ItemMod);
178    
179    // For now, just return the module unchanged
180    // TODO: Implement resource collection and generation
181    quote! {
182        #module
183        
184        // Resources implementation would go here
185    }.into()
186}
187
188/// Mark a module as containing MCP prompts.
189/// Functions in this module will be automatically registered as prompts.
190#[proc_macro_attribute]
191pub fn mcp_prompts(_attr: TokenStream, input: TokenStream) -> TokenStream {
192    let module = parse_macro_input!(input as ItemMod);
193    
194    // For now, just return the module unchanged
195    // TODO: Implement prompt collection and generation
196    quote! {
197        #module
198        
199        // Prompts implementation would go here
200    }.into()
201}
202
203// Keep the original mcp_component for backwards compatibility
204#[proc_macro_attribute]
205pub fn mcp_component(_attr: TokenStream, input: TokenStream) -> TokenStream {
206    mcp_tools(_attr, input)
207}
208
209fn generate_tool_info(func: &ItemFn) -> (TokenStream2, TokenStream2) {
210    let name = func.sig.ident.to_string();
211    let func_ident = &func.sig.ident;
212    let doc = extract_doc_comment(&func.attrs);
213    
214    // Extract parameter information
215    let mut params = Vec::new();
216    let mut param_idents = Vec::new();
217    for input in &func.sig.inputs {
218        if let FnArg::Typed(pat_type) = input {
219            if let Pat::Ident(ident) = &*pat_type.pat {
220                let param_name = ident.ident.to_string();
221                let param_type = type_to_json_schema(&*pat_type.ty);
222                params.push((param_name.clone(), param_type));
223                param_idents.push(ident.ident.clone());
224            }
225        }
226    }
227    
228    // Generate JSON schema for parameters
229    let schema = generate_params_schema(&params);
230    
231    // Generate the tool definition
232    let tool_def = quote! {
233        crate::bindings::fastertools::mcp::tools::Tool {
234            base: crate::bindings::fastertools::mcp::types::BaseMetadata {
235                name: #name.to_string(),
236                title: Some(#name.to_string()),
237            },
238            description: Some(#doc.to_string()),
239            input_schema: #schema.to_string(),
240            output_schema: None,
241            annotations: None,
242            meta: None,
243        }
244    };
245    
246    // Generate the dispatch case
247    let param_extracts: Vec<TokenStream2> = params.iter().zip(&param_idents).map(|((name, _), ident)| {
248        let ident_token = quote! { #ident };
249        quote! {
250            let #ident_token = args[#name].as_str()
251                .ok_or_else(|| crate::bindings::fastertools::mcp::types::McpError {
252                    code: crate::bindings::fastertools::mcp::types::ErrorCode::InvalidParams,
253                    message: format!("Missing or invalid field: {}", #name),
254                    data: None,
255                })?
256                .to_string();
257        }
258    }).collect();
259    
260    let is_async = func.sig.asyncness.is_some();
261    let call_expr = if is_async {
262        quote! { spin_executor::run(#func_ident(#(#param_idents),*)) }
263    } else {
264        quote! { #func_ident(#(#param_idents),*) }
265    };
266    
267    let dispatch = quote! {
268        #name => {
269            #(#param_extracts)*
270            
271            match #call_expr {
272                Ok(result) => Ok(crate::bindings::fastertools::mcp::tools::ToolResult {
273                    content: vec![crate::bindings::fastertools::mcp::types::ContentBlock::Text(
274                        crate::bindings::fastertools::mcp::types::TextContent {
275                            text: result,
276                            annotations: None,
277                            meta: None,
278                        }
279                    )],
280                    is_error: Some(false),
281                    structured_content: None,
282                    meta: None,
283                }),
284                Err(e) => Ok(crate::bindings::fastertools::mcp::tools::ToolResult {
285                    content: vec![crate::bindings::fastertools::mcp::types::ContentBlock::Text(
286                        crate::bindings::fastertools::mcp::types::TextContent {
287                            text: format!("Error: {}", e),
288                            annotations: None,
289                            meta: None,
290                        }
291                    )],
292                    is_error: Some(true),
293                    structured_content: None,
294                    meta: None,
295                })
296            }
297        }
298    };
299    
300    (tool_def, dispatch)
301}
302
303fn extract_doc_comment(attrs: &[Attribute]) -> String {
304    attrs
305        .iter()
306        .filter_map(|attr| {
307            if attr.path().is_ident("doc") {
308                attr.parse_args::<syn::LitStr>().ok().map(|lit| lit.value())
309            } else {
310                None
311            }
312        })
313        .collect::<Vec<_>>()
314        .join(" ")
315        .trim()
316        .to_string()
317}
318
319fn type_to_json_schema(ty: &Type) -> String {
320    match ty {
321        Type::Path(path) => {
322            if let Some(ident) = path.path.get_ident() {
323                match ident.to_string().as_str() {
324                    "i32" | "i64" | "u32" | "u64" => "\"type\": \"integer\"".to_string(),
325                    "f32" | "f64" => "\"type\": \"number\"".to_string(),
326                    "String" => "\"type\": \"string\"".to_string(),
327                    "bool" => "\"type\": \"boolean\"".to_string(),
328                    _ => "\"type\": \"object\"".to_string(),
329                }
330            } else {
331                "\"type\": \"object\"".to_string()
332            }
333        }
334        _ => "\"type\": \"object\"".to_string(),
335    }
336}
337
338fn generate_params_schema(params: &[(String, String)]) -> String {
339    if params.is_empty() {
340        return "{}".to_string();
341    }
342    
343    let properties: Vec<String> = params
344        .iter()
345        .map(|(name, schema)| format!("\"{}\":  {{{}}}", name, schema))
346        .collect();
347    
348    let required: Vec<String> = params
349        .iter()
350        .map(|(name, _)| format!("\"{}\"", name))
351        .collect();
352    
353    format!(
354        "{{\"type\": \"object\", \"properties\": {{{}}}, \"required\": [{}]}}",
355        properties.join(", "),
356        required.join(", ")
357    )
358}
359
360/// Mark a function as an MCP tool (for backwards compatibility).
361#[proc_macro_attribute]
362pub fn tool(_attr: TokenStream, input: TokenStream) -> TokenStream {
363    input
364}
365
366/// Mark a function as an MCP resource.
367#[proc_macro_attribute]
368pub fn resource(_attr: TokenStream, input: TokenStream) -> TokenStream {
369    input
370}
371
372/// Mark a function as an MCP prompt.
373#[proc_macro_attribute]
374pub fn prompt(_attr: TokenStream, input: TokenStream) -> TokenStream {
375    input
376}
wasmcp_macros/lib.rs

wasmcp_macros/
lib.rs
