1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, FnArg, ItemFn};
4
5#[proc_macro_attribute]
12pub fn tool(args: TokenStream, input: TokenStream) -> TokenStream {
13 let input_fn = parse_macro_input!(input as ItemFn);
14
15 let args_parsed = match syn::parse::<ToolArgs>(args) {
17 Ok(args) => args,
18 Err(e) => return e.to_compile_error().into(),
19 };
20
21 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 let doc_comment = extract_doc_comment(&input_fn.attrs);
40
41 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 let input_schema = match args_parsed.input_schema {
62 Some(schema) => schema,
63 None => {
64 quote!(::serde_json::to_value(::schemars::schema_for!(#input_type)).unwrap())
66 }
67 };
68
69 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 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 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 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
157struct ToolArgs {
159 name: Option<String>,
160 title: Option<String>,
161 description: Option<String>,
162 input_schema: Option<proc_macro2::TokenStream>,
163}
164
165fn 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 return Some(doc.trim_start_matches(' ').to_string());
177 }
178 }
179 }
180 }
181 None
182 })
183 .next()
184}
185
186fn 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 Ok(ToolArgs {
244 name,
245 title,
246 description,
247 input_schema,
248 })
249 }
250}