use proc_macro2::TokenStream;
use quote::quote;
use syn::{parse2, Error, ImplItemFn, Result};
use crate::attrs::PromptAttrs;
pub fn expand_prompt(attr: TokenStream, item: TokenStream) -> Result<TokenStream> {
let attrs =
PromptAttrs::parse(attr).map_err(|e| Error::new(proc_macro2::Span::call_site(), e))?;
let method: ImplItemFn = parse2(item)?;
validate_prompt_method(&method)?;
let description = &attrs.description;
let prompt_name = attrs.name.unwrap_or_else(|| method.sig.ident.to_string());
let marker_name = syn::Ident::new(
&format!("__MCP_PROMPT_{}", method.sig.ident),
method.sig.ident.span(),
);
Ok(quote! {
#[doc(hidden)]
#[allow(non_upper_case_globals)]
const #marker_name: (&str, &str) = (#prompt_name, #description);
#[doc = #description]
#[allow(dead_code)]
#method
})
}
fn validate_prompt_method(method: &ImplItemFn) -> Result<()> {
if method.sig.receiver().is_none() {
return Err(Error::new_spanned(
&method.sig,
"prompt methods must take &self",
));
}
if let Some(receiver) = method.sig.receiver() {
if receiver.mutability.is_some() {
return Err(Error::new_spanned(
receiver,
"prompt methods should take &self, not &mut self\n\
help: use interior mutability (e.g., Mutex, RwLock) if you need to modify state",
));
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_validate_prompt_method() {
let method: ImplItemFn = parse_quote! {
async fn greeting(&self, name: String) -> GetPromptResult {
GetPromptResult {
description: None,
messages: vec![],
}
}
};
assert!(validate_prompt_method(&method).is_ok());
let method: ImplItemFn = parse_quote! {
async fn greeting(name: String) -> GetPromptResult {
GetPromptResult {
description: None,
messages: vec![],
}
}
};
assert!(validate_prompt_method(&method).is_err());
}
}