mcpkit-macros 0.2.4

//! Implementation of the `#[prompt]` attribute macro.
//!
//! This module handles parsing and code generation for prompt methods.
//!
//! The main entry point is [`expand_prompt`], which transforms methods annotated
//! with `#[prompt]` into methods with metadata markers that `#[mcp_server]` can discover.

use proc_macro2::TokenStream;
use quote::quote;
use syn::{Error, ImplItemFn, Result, parse2};

use crate::attrs::PromptAttrs;

/// Expand the `#[prompt]` attribute.
///
/// When used standalone (not inside `#[mcp_server]`), this macro
/// preserves the method but adds metadata that `#[mcp_server]` can discover.
pub fn expand_prompt(attr: TokenStream, item: TokenStream) -> Result<TokenStream> {
    // Parse the attribute
    let attrs =
        PromptAttrs::parse(attr).map_err(|e| Error::new(proc_macro2::Span::call_site(), e))?;

    // Parse the method
    let method: ImplItemFn = parse2(item)?;

    // Validate the method signature
    validate_prompt_method(&method)?;

    // Extract info
    let description = &attrs.description;
    let prompt_name = attrs.name.unwrap_or_else(|| method.sig.ident.to_string());

    // Generate a hidden constant that mcp_server can find
    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
    })
}

/// Validate that a method has a valid signature for a prompt.
fn validate_prompt_method(method: &ImplItemFn) -> Result<()> {
    // Must have &self receiver
    if method.sig.receiver().is_none() {
        return Err(Error::new_spanned(
            &method.sig,
            "prompt methods must take &self",
        ));
    }

    // Check that receiver is &self (not &mut self or 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() {
        // Valid 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());

        // Method without self - invalid
        let method: ImplItemFn = parse_quote! {
            async fn greeting(name: String) -> GetPromptResult {
                GetPromptResult {
                    description: None,
                    messages: vec![],
                }
            }
        };
        assert!(validate_prompt_method(&method).is_err());
    }
}