mcp-host-macros 0.1.3

//! #[mcp_tool] attribute macro implementation
//!
//! Transforms async functions into MCP tool handlers with metadata generation.

use darling::{ast::NestedMeta, FromMeta};
use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::{format_ident, quote};
use syn::{parse_macro_input, FnArg, ItemFn, Type};

/// Task support level for tools
#[derive(Debug, Clone, Copy, Default, FromMeta)]
pub enum TaskSupportAttr {
    /// Tool does not support task-augmented execution
    #[default]
    Forbidden,
    /// Tool may support task-augmented execution
    Optional,
    /// Tool requires task-augmented execution
    Required,
}

/// Attributes for #[mcp_tool]
#[derive(Debug, FromMeta)]
pub struct McpToolAttrs {
    /// Tool name (defaults to function name)
    #[darling(default)]
    pub name: Option<String>,

    /// Tool description (defaults to doc comments)
    #[darling(default)]
    pub description: Option<String>,

    /// Visibility predicate expression as string
    #[darling(default)]
    pub visible: Option<String>,

    /// Task support level (forbidden, optional, required)
    #[darling(default)]
    pub task_support: Option<TaskSupportAttr>,

    /// Whether the tool is read-only (hint for LLMs)
    #[darling(default)]
    pub read_only: Option<bool>,

    /// Whether the tool is destructive (hint for LLMs)
    #[darling(default)]
    pub destructive: Option<bool>,

    /// Whether the tool is idempotent (hint for LLMs)
    #[darling(default)]
    pub idempotent: Option<bool>,
}

/// Parse and transform a function marked with #[mcp_tool]
pub fn expand_mcp_tool(attr: TokenStream, item: TokenStream) -> TokenStream {
    let attr_args = match NestedMeta::parse_meta_list(attr.into()) {
        Ok(v) => v,
        Err(e) => return TokenStream::from(e.to_compile_error()),
    };

    let attrs = match McpToolAttrs::from_list(&attr_args) {
        Ok(v) => v,
        Err(e) => return TokenStream::from(e.write_errors()),
    };

    let input_fn = parse_macro_input!(item as ItemFn);

    match generate_tool_impl(attrs, input_fn) {
        Ok(tokens) => tokens.into(),
        Err(e) => e.to_compile_error().into(),
    }
}

fn generate_tool_impl(attrs: McpToolAttrs, input_fn: ItemFn) -> syn::Result<TokenStream2> {
    let fn_name = &input_fn.sig.ident;
    let fn_name_str = fn_name.to_string();
    let tool_name = attrs.name.clone().unwrap_or_else(|| fn_name_str.clone());

    // Generate info function name
    let info_fn_name = format_ident!("{}_tool_info", fn_name);
    let handler_fn_name = format_ident!("{}_handler", fn_name);

    // Extract Parameters<T> type from function signature
    let params_type = extract_parameters_type(&input_fn)?;

    // Generate description
    let description_tokens = match &attrs.description {
        Some(desc) => quote! { Some(#desc.to_string()) },
        None => {
            // Try to extract from doc comments
            let doc = extract_doc_comment(&input_fn);
            if let Some(doc) = doc {
                quote! { Some(#doc.to_string()) }
            } else {
                quote! { None }
            }
        }
    };

    // Generate execution metadata (task support)
    let execution_tokens = match attrs.task_support {
        Some(TaskSupportAttr::Required) => quote! {
            Some(mcp_host::protocol::types::ToolExecution {
                task_support: Some(mcp_host::protocol::types::TaskSupport::Required),
            })
        },
        Some(TaskSupportAttr::Optional) => quote! {
            Some(mcp_host::protocol::types::ToolExecution {
                task_support: Some(mcp_host::protocol::types::TaskSupport::Optional),
            })
        },
        Some(TaskSupportAttr::Forbidden) | None => quote! { None },
    };

    // Generate visibility function if specified
    let visibility_fn = if let Some(visible_expr) = &attrs.visible {
        let vis_fn_name = format_ident!("{}_visibility", fn_name);
        // Parse the expression string into tokens
        let expr: syn::Expr = syn::parse_str(visible_expr).map_err(|e| {
            syn::Error::new_spanned(
                &input_fn.sig.ident,
                format!("Invalid visibility expression: {}", e),
            )
        })?;

        Some(quote! {
            /// Generated visibility predicate for this tool
            pub fn #vis_fn_name(ctx: &mcp_host::server::visibility::VisibilityContext) -> bool {
                #expr
            }
        })
    } else {
        None
    };

    // Generate annotations comment for hints (for documentation)
    let hints_comment = generate_hints_comment(&attrs);

    // Generate the output
    let expanded = quote! {
        #hints_comment
        #input_fn

        /// Generated tool metadata
        pub fn #info_fn_name() -> mcp_host::registry::tools::ToolInfo {
            mcp_host::registry::tools::ToolInfo {
                name: #tool_name.to_string(),
                description: #description_tokens,
                input_schema: mcp_host::macros::schema_for::<#params_type>(),
                execution: #execution_tokens,
            }
        }

        /// Generated handler wrapper
        pub fn #handler_fn_name<'a>(
            server: &'a Self,
            ctx: mcp_host::macros::Ctx<'a>,
            params: serde_json::Value,
        ) -> mcp_host::registry::router::ToolHandlerFuture<'a> {
            Box::pin(async move {
                let params: mcp_host::macros::Parameters<#params_type> = serde_json::from_value(params)
                    .map_err(|e| mcp_host::registry::tools::ToolError::InvalidArguments(e.to_string()))?;
                server.#fn_name(ctx, params).await
            })
        }

        #visibility_fn
    };

    Ok(expanded)
}

/// Generate a doc comment with hints for LLMs (not yet part of protocol, but useful for docs)
fn generate_hints_comment(attrs: &McpToolAttrs) -> TokenStream2 {
    let mut hints = Vec::new();

    if let Some(true) = attrs.read_only {
        hints.push("read-only");
    }
    if let Some(true) = attrs.destructive {
        hints.push("destructive");
    }
    if let Some(true) = attrs.idempotent {
        hints.push("idempotent");
    }

    if hints.is_empty() {
        quote! {}
    } else {
        let hint_str = hints.join(", ");
        quote! {
            #[doc = ""]
            #[doc = concat!("**Hints:** ", #hint_str)]
        }
    }
}

/// Extract the inner type T from Parameters<T> in the function signature
fn extract_parameters_type(input_fn: &ItemFn) -> syn::Result<Type> {
    for arg in &input_fn.sig.inputs {
        if let FnArg::Typed(pat_type) = arg
            && let Type::Path(type_path) = &*pat_type.ty
            && let Some(segment) = type_path.path.segments.last()
            && segment.ident == "Parameters"
            && let syn::PathArguments::AngleBracketed(args) = &segment.arguments
            && let Some(syn::GenericArgument::Type(inner_type)) = args.args.first()
        {
            return Ok(inner_type.clone());
        }
    }

    // If no Parameters<T> found, return empty object type
    // This allows tools with no parameters
    Ok(syn::parse_quote!(()))
}

/// Extract doc comment from function attributes
fn extract_doc_comment(input_fn: &ItemFn) -> Option<String> {
    let mut docs = Vec::new();

    for attr in &input_fn.attrs {
        if attr.path().is_ident("doc")
            && let syn::Meta::NameValue(nv) = &attr.meta
            && let syn::Expr::Lit(expr_lit) = &nv.value
            && let syn::Lit::Str(lit_str) = &expr_lit.lit
        {
            docs.push(lit_str.value().trim().to_string());
        }
    }

    if docs.is_empty() {
        None
    } else {
        Some(docs.join(" ").trim().to_string())
    }
}