use darling::{FromMeta, ast::NestedMeta}; use proc_macro2::{Span, TokenStream}; use quote::{format_ident, quote}; use syn::{Expr, Ident, ImplItemFn, ReturnType}; use crate::common::{extract_doc_line, none_expr}; #[derive(FromMeta, Default, Debug)] #[darling(default)] pub struct PromptAttribute { /// The name of the prompt pub name: Option, /// Human readable title of prompt pub title: Option, /// Optional description of what the prompt does pub description: Option, /// Arguments that can be passed to the prompt pub arguments: Option, /// Optional icons for the prompt pub icons: Option, /// Optional metadata for the prompt pub meta: Option, /// When true, the generated future will not require `Send`. Useful for `!Send` handlers /// (e.g. single-threaded database connections). Also enabled globally by the `local` crate feature. pub local: bool, } pub struct ResolvedPromptAttribute { pub name: String, pub title: Option, pub description: Option, pub arguments: Expr, pub icons: Option, pub meta: Option, } impl ResolvedPromptAttribute { pub fn into_fn(self, fn_ident: Ident) -> syn::Result { let Self { name, description, arguments, title, icons, meta, } = self; let description = if let Some(description) = description { quote! { Some::(#description.into()) } } else { quote! { None:: } }; let title_call = title .map(|t| quote! { .with_title(#t) }) .unwrap_or_default(); let icons_call = icons .map(|i| quote! { .with_icons(#i) }) .unwrap_or_default(); let meta_call = meta.map(|m| quote! { .with_meta(#m) }).unwrap_or_default(); let tokens = quote! { pub fn #fn_ident() -> rmcp::model::Prompt { rmcp::model::Prompt::from_raw( #name, #description, #arguments, ) #title_call #icons_call #meta_call } }; syn::parse2::(tokens) } } pub fn prompt(attr: TokenStream, input: TokenStream) -> syn::Result { let attribute = if attr.is_empty() { Default::default() } else { let attr_args = NestedMeta::parse_meta_list(attr)?; PromptAttribute::from_list(&attr_args)? }; let mut fn_item = syn::parse2::(input.clone())?; let fn_ident = &fn_item.sig.ident; let omit_send = cfg!(feature = "local") || attribute.local; let prompt_attr_fn_ident = format_ident!("{}_prompt_attr", fn_ident); // Try to find prompt parameters from function parameters let arguments_expr = if let Some(arguments) = attribute.arguments { arguments } else { // Look for a type named Parameters in the function signature let params_ty = crate::common::find_parameters_type_impl(&fn_item); if let Some(params_ty) = params_ty { // Generate arguments from the type's schema with caching syn::parse2::(quote! { rmcp::handler::server::prompt::cached_arguments_from_schema::<#params_ty>() })? } else { // No arguments none_expr()? } }; let name = attribute.name.unwrap_or_else(|| fn_ident.to_string()); let description = if let Some(s) = attribute.description { Some(Expr::Lit(syn::ExprLit { attrs: Vec::new(), lit: syn::Lit::Str(syn::LitStr::new(&s, Span::call_site())), })) } else { fn_item.attrs.iter().try_fold(None, extract_doc_line)? }; let arguments = arguments_expr; let resolved_prompt_attr = ResolvedPromptAttribute { name: name.clone(), description: description.clone(), arguments: arguments.clone(), title: attribute.title, icons: attribute.icons, meta: attribute.meta, }; let prompt_attr_fn = resolved_prompt_attr.into_fn(prompt_attr_fn_ident.clone())?; // Modify the input function for async support (same as tool macro) if fn_item.sig.asyncness.is_some() { // 1. remove asyncness from sig // 2. make return type: `std::pin::Pin + Send + '_>>` // (omit `+ Send` when the `local` crate feature is active or `#[prompt(local)]` is used) // 3. make body: { Box::pin(async move { #body }) } let new_output = syn::parse2::({ let mut lt = quote! { 'static }; if let Some(receiver) = fn_item.sig.receiver() && let syn::ReceiverKind::Reference(_, receiver_lt, _) = &receiver.kind { if let Some(receiver_lt) = receiver_lt { lt = quote! { #receiver_lt }; } else { lt = quote! { '_ }; } } match &fn_item.sig.output { syn::ReturnType::Default => { if omit_send { quote! { -> ::std::pin::Pin + #lt>> } } else { quote! { -> ::std::pin::Pin + Send + #lt>> } } } syn::ReturnType::Type(_, ty) => { if omit_send { quote! { -> ::std::pin::Pin + #lt>> } } else { quote! { -> ::std::pin::Pin + Send + #lt>> } } } } })?; let prev_block = &fn_item.block; let new_block = syn::parse2::(quote! { { Box::pin(async move #prev_block ) } })?; fn_item.sig.asyncness = None; fn_item.sig.output = new_output; fn_item.block = new_block; } Ok(quote! { #prompt_attr_fn #fn_item }) } #[cfg(test)] mod test { use super::*; #[test] fn test_prompt_macro() -> syn::Result<()> { let attr = quote! { name = "example-prompt", description = "An example prompt" }; let input = quote! { async fn example_prompt(&self, Parameters(args): Parameters) -> Result { Ok("Example prompt response".to_string()) } }; let result = prompt(attr, input)?; // Verify the output contains both the attribute function and the modified function let result_str = result.to_string(); assert!(result_str.contains("example_prompt_prompt_attr")); assert!( result_str.contains("rmcp") && result_str.contains("model") && result_str.contains("Prompt") ); Ok(()) } #[test] fn test_doc_comment_description() -> syn::Result<()> { let attr = quote! {}; // No explicit description let input = quote! { /// This is a test prompt description /// with multiple lines fn test_prompt(&self) -> Result { Ok("Test".to_string()) } }; let result = prompt(attr, input)?; // The output should contain the description from doc comments let result_str = result.to_string(); assert!(result_str.contains("This is a test prompt description")); assert!(result_str.contains("with multiple lines")); Ok(()) } #[test] fn test_doc_include_description() -> syn::Result<()> { let attr = quote! {}; // No explicit description let input = quote! { #[doc = include_str!("some/test/data/doc.txt")] fn test_prompt_included(&self) -> Result { Ok("Test".to_string()) } }; let result = prompt(attr, input)?; // The generated tokens should preserve the include_str! invocation let result_str = result.to_string(); assert!(result_str.contains("include_str")); Ok(()) } #[test] fn test_async_prompt_default_send_behavior() -> syn::Result<()> { let attr = quote! {}; let input = quote! { async fn test_prompt_default_send(&self) -> String { "ok".to_string() } }; let result = prompt(attr, input)?; let result_str = result.to_string(); if cfg!(feature = "local") { assert!(!result_str.contains("+ Send +")); } else { assert!(result_str.contains("+ Send +")); } Ok(()) } #[test] fn test_async_prompt_preserves_receiver_lifetime() -> syn::Result<()> { let attr = quote! {}; let input = quote! { async fn test_prompt_with_lifetime<'a>(&'a self) -> String { "ok".to_string() } }; let result = prompt(attr, input)?; assert!(result.to_string().contains("+ 'a")); Ok(()) } #[test] fn test_async_prompt_local_omits_send() -> syn::Result<()> { let attr = quote! { local }; let input = quote! { async fn test_prompt_local_no_send(&self) -> String { "ok".to_string() } }; let result = prompt(attr, input)?; let result_str = result.to_string(); assert!(!result_str.contains("+ Send +")); Ok(()) } }