ftl_sdk_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, FnArg, ItemFn};
4
5/// Create a tool that can be used with the MCP Gateway.
6///
7/// The macro will:
8/// - Use the function name as the tool name (unless overridden)
9/// - Extract the first line of the doc comment as the description (unless overridden)
10/// - Generate the title from the function name (unless overridden)
11#[proc_macro_attribute]
12pub fn tool(args: TokenStream, input: TokenStream) -> TokenStream {
13    let input_fn = parse_macro_input!(input as ItemFn);
14
15    // Parse the arguments to extract name, title, description
16    let args_parsed = match syn::parse::<ToolArgs>(args) {
17        Ok(args) => args,
18        Err(e) => return e.to_compile_error().into(),
19    };
20
21    // Get the input type to derive the schema
22    let input_type = match input_fn.sig.inputs.first() {
23        Some(FnArg::Typed(pat_type)) => &pat_type.ty,
24        _ => {
25            return syn::Error::new_spanned(
26                &input_fn.sig,
27                "Function must have exactly one argument",
28            )
29            .to_compile_error()
30            .into();
31        }
32    };
33
34    let fn_name = &input_fn.sig.ident;
35    let fn_visibility = &input_fn.vis;
36    let is_async = input_fn.sig.asyncness.is_some();
37
38    // Extract doc comments from the function
39    let doc_comment = extract_doc_comment(&input_fn.attrs);
40
41    // Use provided values or fall back to defaults
42    let name = args_parsed.name.unwrap_or_else(|| fn_name.to_string());
43    let title = args_parsed
44        .title
45        .map(|s| quote!(Some(#s.to_string())))
46        .unwrap_or_else(|| {
47            let title = generate_title(&fn_name.to_string());
48            quote!(Some(#title.to_string()))
49        });
50    let description = args_parsed
51        .description
52        .map(|s| quote!(Some(#s.to_string())))
53        .unwrap_or_else(|| {
54            if let Some(doc) = doc_comment {
55                quote!(Some(#doc.to_string()))
56            } else {
57                quote!(None)
58            }
59        });
60    // Generate input schema - either from provided value or derive from type
61    let input_schema = match args_parsed.input_schema {
62        Some(schema) => schema,
63        None => {
64            // Automatically derive schema from the input type
65            quote!(::serde_json::to_value(::schemars::schema_for!(#input_type)).unwrap())
66        }
67    };
68
69    // Generate the function call with or without await
70    let fn_call = if is_async {
71        quote!(#fn_name(input).await)
72    } else {
73        quote!(#fn_name(input))
74    };
75
76    let output = quote! {
77        #input_fn
78
79        #[::spin_sdk::http_component]
80        #fn_visibility async fn handle_tool_component(req: ::spin_sdk::http::Request) -> ::spin_sdk::http::Response {
81            use ::spin_sdk::http::{Method, Response};
82
83            // Build metadata
84            let metadata = ::ftl_sdk::ToolMetadata {
85                name: #name.to_string(),
86                title: #title,
87                description: #description,
88                input_schema: #input_schema,
89                output_schema: None,
90                annotations: None,
91                meta: None,
92            };
93
94            match req.method() {
95                &Method::Get => {
96                    // Return tool metadata
97                    match ::serde_json::to_vec(&metadata) {
98                        Ok(body) => Response::builder()
99                            .status(200)
100                            .header("Content-Type", "application/json")
101                            .body(body)
102                            .build(),
103                        Err(e) => Response::builder()
104                            .status(500)
105                            .body(format!("Failed to serialize metadata: {}", e))
106                            .build()
107                    }
108                }
109                &Method::Post => {
110                    // Parse request body and execute tool
111                    let body = req.body();
112                    match ::serde_json::from_slice::<#input_type>(body) {
113                        Ok(input) => {
114                            let response = #fn_call;
115                            match ::serde_json::to_vec(&response) {
116                                Ok(body) => Response::builder()
117                                    .status(200)
118                                    .header("Content-Type", "application/json")
119                                    .body(body)
120                                    .build(),
121                                Err(e) => {
122                                    let error_response = ::ftl_sdk::ToolResponse::error(
123                                        format!("Failed to serialize response: {}", e)
124                                    );
125                                    Response::builder()
126                                        .status(500)
127                                        .header("Content-Type", "application/json")
128                                        .body(::serde_json::to_vec(&error_response).unwrap_or_default())
129                                        .build()
130                                }
131                            }
132                        }
133                        Err(e) => {
134                            let error_response = ::ftl_sdk::ToolResponse::error(
135                                format!("Invalid request body: {}", e)
136                            );
137                            Response::builder()
138                                .status(400)
139                                .header("Content-Type", "application/json")
140                                .body(::serde_json::to_vec(&error_response).unwrap_or_default())
141                                .build()
142                        }
143                    }
144                }
145                _ => Response::builder()
146                    .status(405)
147                    .header("Allow", "GET, POST")
148                    .body("Method not allowed")
149                    .build()
150            }
151        }
152    };
153
154    output.into()
155}
156
157// Helper struct to parse tool macro arguments
158struct ToolArgs {
159    name: Option<String>,
160    title: Option<String>,
161    description: Option<String>,
162    input_schema: Option<proc_macro2::TokenStream>,
163}
164
165// Extract the first line of doc comments from attributes
166fn extract_doc_comment(attrs: &[syn::Attribute]) -> Option<String> {
167    attrs
168        .iter()
169        .filter_map(|attr| {
170            if attr.path().is_ident("doc") {
171                if let syn::Meta::NameValue(nv) = &attr.meta {
172                    if let syn::Expr::Lit(lit) = &nv.value {
173                        if let syn::Lit::Str(s) = &lit.lit {
174                            let doc = s.value();
175                            // Trim leading space that rustdoc adds
176                            return Some(doc.trim_start_matches(' ').to_string());
177                        }
178                    }
179                }
180            }
181            None
182        })
183        .next()
184}
185
186// Generate a title from a function name (e.g., "calculate_sum" -> "Calculate Sum")
187fn generate_title(name: &str) -> String {
188    name.split('_')
189        .map(|word| {
190            let mut chars = word.chars();
191            match chars.next() {
192                None => String::new(),
193                Some(first) => first.to_uppercase().collect::<String>() + chars.as_str(),
194            }
195        })
196        .collect::<Vec<_>>()
197        .join(" ")
198}
199
200impl syn::parse::Parse for ToolArgs {
201    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
202        let mut name = None;
203        let mut title = None;
204        let mut description = None;
205        let mut input_schema = None;
206
207        while !input.is_empty() {
208            let ident: syn::Ident = input.parse()?;
209            input.parse::<syn::Token![=]>()?;
210
211            match ident.to_string().as_str() {
212                "name" => {
213                    let lit: syn::LitStr = input.parse()?;
214                    name = Some(lit.value());
215                }
216                "title" => {
217                    let lit: syn::LitStr = input.parse()?;
218                    title = Some(lit.value());
219                }
220                "description" => {
221                    let lit: syn::LitStr = input.parse()?;
222                    description = Some(lit.value());
223                }
224                "input_schema" => {
225                    let expr: syn::Expr = input.parse()?;
226                    input_schema = Some(quote!(#expr));
227                }
228                _ => {
229                    return Err(syn::Error::new_spanned(
230                        ident,
231                        "Unknown attribute. Expected: name, title, description, or input_schema",
232                    ));
233                }
234            }
235
236            if !input.is_empty() {
237                input.parse::<syn::Token![,]>()?;
238            }
239        }
240
241        // input_schema is now optional
242
243        Ok(ToolArgs {
244            name,
245            title,
246            description,
247            input_schema,
248        })
249    }
250}