server-less-macros 0.6.0

//! MCP (Model Context Protocol) tool generation macro.
//!
//! Generates MCP tool definitions from Rust impl blocks for use with Claude and other LLMs.
//!
//! # What is MCP?
//!
//! [Model Context Protocol](https://modelcontextprotocol.io) is an open standard for exposing
//! tools and context to language models. Each method becomes a callable tool with JSON schema
//! for parameters and return values.
//!
//! # Tool Naming
//!
//! - Methods are exposed with their original names (e.g., `read_file`)
//! - Optional namespace prefix: `#[mcp(namespace = "myapp")]` → `myapp_create_user`
//! - Tool names must be valid identifiers (alphanumeric + underscore)
//! - Namespace is added as prefix with underscore separator
//!
//! # Mount Points
//!
//! Methods returning `&T` become tool namespaces (mount points):
//! - Static: `fn users(&self) -> &Users` → tools prefixed `users_*`
//! - Slug: `fn user(&self, id: String) -> &UserService` → tools prefixed `user_*` with `id` merged
//!
//! The mounted type must implement `McpNamespace` (generated by `#[mcp]`).
//!
//! # Parameter Schema
//!
//! Parameters are automatically converted to JSON schema:
//! - `String` → string
//! - `i32`, `u64`, etc. → integer
//! - `f32`, `f64` → number
//! - `bool` → boolean
//! - `Vec<T>`, `[T]` → array
//! - Custom structs → object (requires Serialize/Deserialize)
//! - `Option<T>` → optional parameter (nullable)
//!
//! # Return Types
//!
//! Return values are automatically converted to JSON:
//!
//! - `()` → `{"success": true}`
//! - `T` → serialized value
//! - `Result<T, E>` → `T` on success, error string on failure
//! - `Option<T>` → `T` or `null`
//! - `Vec<T>` → JSON array
//! - `impl Stream<Item = T>` → collected into JSON array
//!
//! # Generated Methods
//!
//! - `mcp_tools() -> Vec<serde_json::Value>` - Tool definitions for MCP
//! - `mcp_method_names() -> Vec<String>` - List of tool/method names
//! - `mcp_call(&self, name: &str, args: Value) -> Result<Value, String>` - Execute tool
//! - `mcp_call_async(&self, name: &str, args: Value).await` - Async execution
//!
//! Also implements `McpNamespace` trait for composition.

use crate::context::partition_context_params;
use crate::server_attrs::{has_server_hidden, has_server_skip, validate_server_attrs};
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use server_less_parse::{MethodInfo, extract_methods, get_impl_name, partition_methods};
use server_less_rpc::{self, AsyncHandling};
use syn::{ItemImpl, Token, parse::Parse};

/// Arguments for the #[mcp] attribute
#[derive(Default)]
pub(crate) struct McpArgs {
    /// Tool namespace/prefix
    pub(crate) namespace: Option<String>,
    /// App name (from `#[app]` or inline; used as namespace fallback)
    pub name: Option<String>,
    /// App description (from `#[app]` or inline)
    pub description: Option<String>,
}

impl Parse for McpArgs {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let mut args = McpArgs::default();

        while !input.is_empty() {
            let ident: syn::Ident = input.parse()?;
            input.parse::<Token![=]>()?;

            match ident.to_string().as_str() {
                "namespace" => {
                    let lit: syn::LitStr = input.parse()?;
                    args.namespace = Some(lit.value());
                }
                other => {
                    const VALID: &[&str] = &["namespace"];
                    let suggestion = crate::did_you_mean(other, VALID)
                        .map(|s| format!(" — did you mean `{s}`?"))
                        .unwrap_or_default();
                    return Err(syn::Error::new(
                        ident.span(),
                        format!(
                            "unknown argument `{other}`{suggestion}. Valid arguments: namespace\n\
                             \n\
                             Related: #[tool] preset (MCP + JSON Schema), #[jsonschema] (standalone schema)"
                        ),
                    ));
                }
            }

            if input.peek(Token![,]) {
                input.parse::<Token![,]>()?;
            }
        }

        Ok(args)
    }
}

/// Strip `#[param(...)]` attributes from function parameters in the re-emitted impl block.
/// These are consumed by the `#[mcp]` macro during parsing; leaving them in the output
/// would cause "cannot find attribute `param`" errors if `#[http]` is not also applied.
fn strip_param_attrs(impl_block: &ItemImpl) -> ItemImpl {
    let mut block = impl_block.clone();
    for item in &mut block.items {
        if let syn::ImplItem::Fn(method) = item {
            for input in &mut method.sig.inputs {
                if let syn::FnArg::Typed(pat_type) = input {
                    pat_type.attrs.retain(|attr| !attr.path().is_ident("param"));
                }
            }
        }
    }
    block
}

pub(crate) fn expand_mcp(args: McpArgs, mut impl_block: ItemImpl) -> syn::Result<TokenStream2> {
    // Extract #[__app_meta] from attrs and use as fallback for unset fields.
    let app_meta = crate::app::extract_app_meta(&mut impl_block.attrs);
    let app_name = args.name.or(app_meta.name);
    let _app_description = args.description.or(app_meta.description);

    crate::reject_generic_impl(&impl_block)?;
    let _struct_name = get_impl_name(&impl_block)?;
    let (impl_generics, _ty_generics, where_clause) = impl_block.generics.split_for_impl();
    let self_ty = &impl_block.self_ty;
    let methods = extract_methods(&impl_block)?;
    let clean_impl = if crate::is_protocol_impl_emitter(&impl_block, "mcp") {
        let stripped = strip_param_attrs(&impl_block);
        quote! { #stripped }
    } else {
        quote! {}
    };

    // Use explicit namespace first; fall back to app_meta.name (from #[app]) as namespace.
    let namespace = args.namespace.or(app_name).unwrap_or_default();
    let namespace_prefix = if namespace.is_empty() {
        String::new()
    } else {
        format!("{}_", namespace)
    };

    for m in &methods {
        validate_server_attrs(m)?;
    }
    let partitioned = partition_methods(&methods, has_server_skip);

    // Separate hidden from visible leaf methods.
    // Hidden methods are still dispatchable but absent from tool listings.
    let visible_leaf: Vec<_> = partitioned
        .leaf
        .iter()
        .copied()
        .filter(|m| !has_server_hidden(m))
        .collect();

    // Generate tool definitions for visible leaf methods only.
    // Each entry is a statement so #[cfg] can guard individual tool insertions.
    let leaf_tool_definitions: Vec<_> = visible_leaf
        .iter()
        .map(|m| {
            let def = generate_tool_definition(&namespace_prefix, m)?;
            let cfg_attrs = &m.cfg_attrs;
            Ok(quote! {
                #(#cfg_attrs)*
                tools.push(#def);
            })
        })
        .collect::<syn::Result<Vec<_>>>()?;

    // Generate dispatch match arms for ALL leaf methods (hidden methods remain callable)
    let leaf_dispatch_sync: Vec<_> = partitioned
        .leaf
        .iter()
        .map(|m| {
            let arm = generate_dispatch_arm_sync(&namespace_prefix, m);
            let cfg_attrs = &m.cfg_attrs;
            quote! {
                #(#cfg_attrs)*
                #arm
            }
        })
        .collect();

    let leaf_dispatch_async: Vec<_> = partitioned
        .leaf
        .iter()
        .map(|m| {
            let arm = generate_dispatch_arm_async(&namespace_prefix, m);
            let cfg_attrs = &m.cfg_attrs;
            quote! {
                #(#cfg_attrs)*
                #arm
            }
        })
        .collect();

    // Tool names for visible leaf methods only.
    // Each entry is a statement so #[cfg] can guard individual name insertions.
    let leaf_tool_names: Vec<_> = visible_leaf
        .iter()
        .map(|m| {
            let name = format!("{}{}", namespace_prefix, m.name);
            let cfg_attrs = &m.cfg_attrs;
            quote! {
                #(#cfg_attrs)*
                names.push(#name.to_string());
            }
        })
        .collect();

    // Generate mount point contributions
    let mount_tools: Vec<_> = partitioned
        .static_mounts
        .iter()
        .chain(partitioned.slug_mounts.iter())
        .map(|m| generate_mount_tools(&namespace_prefix, m))
        .collect::<syn::Result<Vec<_>>>()?;

    let mount_tool_names: Vec<_> = partitioned
        .static_mounts
        .iter()
        .chain(partitioned.slug_mounts.iter())
        .map(|m| generate_mount_tool_names(&namespace_prefix, m))
        .collect::<syn::Result<Vec<_>>>()?;

    let mount_dispatch_sync: Vec<_> = partitioned
        .static_mounts
        .iter()
        .map(|m| generate_static_mount_dispatch(&namespace_prefix, m, AsyncHandling::Error))
        .chain(
            partitioned
                .slug_mounts
                .iter()
                .map(|m| generate_slug_mount_dispatch(&namespace_prefix, m, AsyncHandling::Error)),
        )
        .collect::<syn::Result<Vec<_>>>()?;

    let mount_dispatch_async: Vec<_> = partitioned
        .static_mounts
        .iter()
        .map(|m| generate_static_mount_dispatch(&namespace_prefix, m, AsyncHandling::Await))
        .chain(
            partitioned
                .slug_mounts
                .iter()
                .map(|m| generate_slug_mount_dispatch(&namespace_prefix, m, AsyncHandling::Await)),
        )
        .collect::<syn::Result<Vec<_>>>()?;

    // Build tool documentation (visible methods only)
    let tool_doc_entries: Vec<String> = visible_leaf
        .iter()
        .map(|m| {
            let name = format!("{}{}", namespace_prefix, m.name);
            match &m.docs {
                Some(doc) => format!("- `{name}` — {doc}"),
                None => format!("- `{name}`"),
            }
        })
        .collect();
    let has_mounts = !partitioned.static_mounts.is_empty() || !partitioned.slug_mounts.is_empty();
    let mcp_tools_doc = if tool_doc_entries.is_empty() && !has_mounts {
        "Get the list of available MCP tool definitions.".to_string()
    } else {
        let mount_note = if has_mounts {
            "\n\nAlso includes tools from mounted sub-services."
        } else {
            ""
        };
        format!(
            "Get the list of available MCP tool definitions.\n\n# Tools\n\n{}{}",
            tool_doc_entries.join("\n"),
            mount_note
        )
    };

    Ok(quote! {
        #clean_impl

        impl #impl_generics ::server_less::McpNamespace for #self_ty #where_clause {
            fn mcp_namespace_tools() -> Vec<::server_less::serde_json::Value> {
                Self::mcp_tools()
            }

            fn mcp_namespace_tool_names() -> Vec<String> {
                Self::mcp_method_names()
            }

            fn mcp_namespace_call(
                &self,
                name: &str,
                args: ::server_less::serde_json::Value,
            ) -> ::std::result::Result<::server_less::serde_json::Value, String> {
                self.mcp_call(name, args)
            }

            async fn mcp_namespace_call_async(
                &self,
                name: &str,
                args: ::server_less::serde_json::Value,
            ) -> ::std::result::Result<::server_less::serde_json::Value, String> {
                self.mcp_call_async(name, args).await
            }
        }

        impl #impl_generics #self_ty #where_clause {
            #[doc = #mcp_tools_doc]
            pub fn mcp_tools() -> Vec<::server_less::serde_json::Value> {
                let mut tools = Vec::new();
                #(#leaf_tool_definitions)*
                #(#mount_tools)*
                tools
            }

            /// Get tool/method names
            pub fn mcp_method_names() -> Vec<String> {
                let mut names: Vec<String> = Vec::new();
                #(#leaf_tool_names)*
                #(#mount_tool_names)*
                names
            }

            /// Call an MCP tool by name with JSON arguments (sync version)
            ///
            /// Note: Async methods will return an error. Use `mcp_call_async` for async methods.
            pub fn mcp_call(
                &self,
                name: &str,
                args: ::server_less::serde_json::Value
            ) -> ::std::result::Result<::server_less::serde_json::Value, String> {
                match name {
                    #(#leaf_dispatch_sync)*
                    #(#mount_dispatch_sync)*
                    _ => Err(format!("Unknown tool: {}", name)),
                }
            }

            /// Call an MCP tool (async version)
            ///
            /// Supports both sync and async methods. Async methods are awaited properly.
            pub async fn mcp_call_async(
                &self,
                name: &str,
                args: ::server_less::serde_json::Value
            ) -> ::std::result::Result<::server_less::serde_json::Value, String> {
                match name {
                    #(#leaf_dispatch_async)*
                    #(#mount_dispatch_async)*
                    _ => Err(format!("Unknown tool: {}", name)),
                }
            }
        }
    })
}

/// Generate MCP parameter schema entries, respecting `#[param(name = "...")]` wire-name overrides.
///
/// Unlike `server_less_rpc::generate_param_schema_for`, this uses the wire name (from
/// `#[param(name)]`) as the JSON key when present, falling back to the Rust identifier name.
fn generate_mcp_param_schema(
    params: &[&server_less_parse::ParamInfo],
) -> (Vec<proc_macro2::TokenStream>, Vec<String>) {
    let properties: Vec<_> = params
        .iter()
        .map(|p| {
            let param_name = p.wire_name.clone().unwrap_or_else(|| p.name_str());
            let param_type = server_less_rpc::infer_json_type(&p.ty);
            let description = p
                .help_text
                .clone()
                .unwrap_or_else(|| format!("Parameter: {}", param_name));
            quote! { (#param_name, #param_type, #description) }
        })
        .collect();

    let required: Vec<_> = params
        .iter()
        .filter(|p| !p.is_optional)
        .map(|p| p.wire_name.clone().unwrap_or_else(|| p.name_str()))
        .collect();

    (properties, required)
}

/// Generate an MCP tool definition (JSON schema)
fn generate_tool_definition(
    namespace_prefix: &str,
    method: &MethodInfo,
) -> syn::Result<TokenStream2> {
    let base_name = method.wire_name_or(|n| n);
    let name = format!("{}{}", namespace_prefix, base_name);
    let description = method
        .docs
        .clone()
        .unwrap_or(base_name.clone());

    // Partition out Context parameters — they are injected, not user-visible inputs.
    let (_ctx_param, user_params) =
        partition_context_params(&method.params)?;

    // Generate parameter schema, honoring #[param(name = "...")] wire-name override.
    let (properties, required_params) = generate_mcp_param_schema(&user_params);

    Ok(quote! {
        {
            let mut properties = ::server_less::serde_json::Map::new();
            #(
                {
                    let (name, type_str, desc): (&str, &str, &str) = #properties;
                    properties.insert(name.to_string(), ::server_less::serde_json::json!({
                        "type": type_str,
                        "description": desc
                    }));
                }
            )*

            ::server_less::serde_json::json!({
                "name": #name,
                "description": #description,
                "inputSchema": {
                    "type": "object",
                    "properties": properties,
                    "required": [#(#required_params),*]
                }
            })
        }
    })
}

/// Generate a dispatch match arm for calling a method (sync version)
fn generate_dispatch_arm_sync(
    namespace_prefix: &str,
    method: &MethodInfo,
) -> TokenStream2 {
    let tool_name = format!("{}{}", namespace_prefix, method.name);
    generate_dispatch_arm_with_context(method, Some(&tool_name), AsyncHandling::Error)
}

/// Generate a dispatch match arm for calling a method (async version)
fn generate_dispatch_arm_async(
    namespace_prefix: &str,
    method: &MethodInfo,
) -> TokenStream2 {
    let tool_name = format!("{}{}", namespace_prefix, method.name);
    generate_dispatch_arm_with_context(method, Some(&tool_name), AsyncHandling::Await)
}

/// Generate a dispatch arm that injects Context parameters instead of reading them from JSON.
fn generate_dispatch_arm_with_context(
    method: &MethodInfo,
    tool_name: Option<&str>,
    async_handling: AsyncHandling,
) -> TokenStream2 {
    // Find context parameter indices for injection (per-method detection).
    let injections: Vec<(usize, TokenStream2)> = method
        .params
        .iter()
        .enumerate()
        .filter_map(|(i, p)| {
            if crate::context::should_inject_context(&p.ty, &method.params) {
                Some((i, quote! { ::server_less::Context::default() }))
            } else {
                None
            }
        })
        .collect();

    if injections.is_empty() {
        server_less_rpc::generate_dispatch_arm(method, tool_name, async_handling)
    } else {
        server_less_rpc::generate_dispatch_arm_with_injections(
            method,
            tool_name,
            async_handling,
            &injections,
        )
    }
}

/// Generate code to append mounted tools to the tools list.
fn generate_mount_tools(namespace_prefix: &str, method: &MethodInfo) -> syn::Result<TokenStream2> {
    let mount_name = method.wire_name_or(|n| n);
    let full_prefix = format!("{}{}_{}", namespace_prefix, mount_name, "");
    let inner_ty = method.return_info.reference_inner.as_ref().ok_or_else(|| {
        syn::Error::new_spanned(
            &method.method.sig,
            "BUG: mount method must have a reference return type (&T)",
        )
    })?;
    let is_slug = !method.params.is_empty();

    if is_slug {
        // For slug mounts, add the slug parameter(s) to each tool's inputSchema
        let slug_params: Vec<_> = method.params.iter().collect();
        let slug_properties: Vec<_> = slug_params
            .iter()
            .map(|p| {
                let name = p.name_str();
                let json_type = server_less_rpc::infer_json_type(&p.ty);
                quote! { (#name, #json_type) }
            })
            .collect();
        let slug_required: Vec<_> = slug_params
            .iter()
            .filter(|p| !p.is_optional)
            .map(|p| p.name_str())
            .collect();

        Ok(quote! {
            {
                let child_tools = <#inner_ty as ::server_less::McpNamespace>::mcp_namespace_tools();
                for mut tool in child_tools {
                    // Prefix the tool name
                    if let Some(name) = tool.get("name").and_then(|n| n.as_str()) {
                        let prefixed = format!("{}{}", #full_prefix, name);
                        tool.as_object_mut().unwrap().insert("name".to_string(),
                            ::server_less::serde_json::Value::String(prefixed));
                    }
                    // Add slug parameters to inputSchema
                    if let Some(schema) = tool.get_mut("inputSchema") {
                        if let Some(props) = schema.get_mut("properties") {
                            if let Some(props_map) = props.as_object_mut() {
                                #(
                                    {
                                        let (slug_name, slug_type): (&str, &str) = #slug_properties;
                                        props_map.insert(slug_name.to_string(),
                                            ::server_less::serde_json::json!({"type": slug_type, "description": format!("Parameter: {}", slug_name)}));
                                    }
                                )*
                            }
                        }
                        if let Some(required) = schema.get_mut("required") {
                            if let Some(req_arr) = required.as_array_mut() {
                                #(
                                    req_arr.push(::server_less::serde_json::Value::String(#slug_required.to_string()));
                                )*
                            }
                        }
                    }
                    tools.push(tool);
                }
            }
        })
    } else {
        // Static mount: just prefix the names
        Ok(quote! {
            {
                let child_tools = <#inner_ty as ::server_less::McpNamespace>::mcp_namespace_tools();
                for mut tool in child_tools {
                    if let Some(name) = tool.get("name").and_then(|n| n.as_str()) {
                        let prefixed = format!("{}{}", #full_prefix, name);
                        tool.as_object_mut().unwrap().insert("name".to_string(),
                            ::server_less::serde_json::Value::String(prefixed));
                    }
                    tools.push(tool);
                }
            }
        })
    }
}

/// Generate code to append mounted tool names to the names list.
fn generate_mount_tool_names(
    namespace_prefix: &str,
    method: &MethodInfo,
) -> syn::Result<TokenStream2> {
    let mount_name = method.wire_name_or(|n| n);
    let full_prefix = format!("{}{}_{}", namespace_prefix, mount_name, "");
    let inner_ty = method.return_info.reference_inner.as_ref().ok_or_else(|| {
        syn::Error::new_spanned(
            &method.method.sig,
            "BUG: mount method must have a reference return type (&T)",
        )
    })?;

    Ok(quote! {
        {
            let child_names = <#inner_ty as ::server_less::McpNamespace>::mcp_namespace_tool_names();
            for child_name in child_names {
                let prefixed = format!("{}{}", #full_prefix, child_name);
                names.push(prefixed);
            }
        }
    })
}

/// Generate dispatch for a static mount (`fn foo(&self) -> &T`).
fn generate_static_mount_dispatch(
    namespace_prefix: &str,
    method: &MethodInfo,
    async_handling: AsyncHandling,
) -> syn::Result<TokenStream2> {
    let mount_name = method.wire_name_or(|n| n);
    let mount_prefix = format!("{}{}_{}", namespace_prefix, mount_name, "");
    let method_name = &method.name;
    let inner_ty = method.return_info.reference_inner.as_ref().ok_or_else(|| {
        syn::Error::new_spanned(
            &method.method.sig,
            "BUG: mount method must have a reference return type (&T)",
        )
    })?;

    Ok(match async_handling {
        AsyncHandling::Await => quote! {
            __name if __name.starts_with(#mount_prefix) => {
                let __stripped = &__name[#mount_prefix.len()..];
                let __delegate = self.#method_name();
                <#inner_ty as ::server_less::McpNamespace>::mcp_namespace_call_async(__delegate, __stripped, args).await
            }
        },
        _ => quote! {
            __name if __name.starts_with(#mount_prefix) => {
                let __stripped = &__name[#mount_prefix.len()..];
                let __delegate = self.#method_name();
                <#inner_ty as ::server_less::McpNamespace>::mcp_namespace_call(__delegate, __stripped, args)
            }
        },
    })
}

/// Generate dispatch for a slug mount (`fn foo(&self, id: Id) -> &T`).
fn generate_slug_mount_dispatch(
    namespace_prefix: &str,
    method: &MethodInfo,
    async_handling: AsyncHandling,
) -> syn::Result<TokenStream2> {
    let mount_name = method.wire_name_or(|n| n);
    let mount_prefix = format!("{}{}_{}", namespace_prefix, mount_name, "");
    let method_name = &method.name;
    let inner_ty = method.return_info.reference_inner.as_ref().ok_or_else(|| {
        syn::Error::new_spanned(
            &method.method.sig,
            "BUG: mount method must have a reference return type (&T)",
        )
    })?;

    // Generate slug parameter extraction from args
    let slug_extractions: Vec<_> = method
        .params
        .iter()
        .map(server_less_rpc::generate_param_extraction)
        .collect();
    let slug_names: Vec<_> = method.params.iter().map(|p| &p.name).collect();

    Ok(match async_handling {
        AsyncHandling::Await => quote! {
            __name if __name.starts_with(#mount_prefix) => {
                let __stripped = &__name[#mount_prefix.len()..];
                #(#slug_extractions)*
                let __delegate = self.#method_name(#(#slug_names),*);
                <#inner_ty as ::server_less::McpNamespace>::mcp_namespace_call_async(__delegate, __stripped, args).await
            }
        },
        _ => quote! {
            __name if __name.starts_with(#mount_prefix) => {
                let __stripped = &__name[#mount_prefix.len()..];
                #(#slug_extractions)*
                let __delegate = self.#method_name(#(#slug_names),*);
                <#inner_ty as ::server_less::McpNamespace>::mcp_namespace_call(__delegate, __stripped, args)
            }
        },
    })
}