klieo-macros 0.8.0

//! Generate the TokenStream for `#[tool]`.

use crate::parse::ToolDecl;
use proc_macro2::TokenStream;
use quote::{format_ident, quote};

pub(crate) fn expand(decl: ToolDecl) -> TokenStream {
    let ToolDecl { attr, func } = decl;

    let original_fn = &func.item_fn;
    let fn_name = &func.fn_name;
    let fn_name_str = fn_name.to_string();

    // Tool name shown to the LLM: `rename` overrides the snake_case fn name.
    let tool_name_str = match &attr.rename {
        Some(r) => r.value(),
        None => fn_name_str.clone(),
    };

    // Description: explicit attribute → first doc-comment line → compile error.
    let description_str = match &attr.description {
        Some(d) => d.value(),
        None => match crate::parse::first_doc_line(&func.item_fn.attrs) {
            Some(s) => s,
            None => {
                return syn::Error::new_spanned(
                    fn_name,
                    "missing tool description: supply `#[tool(description = \"…\")]` or a `///` doc comment on the function",
                )
                .to_compile_error();
            }
        },
    };

    // PascalCase + "Tool" suffix for the public struct.
    let struct_name = format_ident!("{}Tool", to_pascal_case(&fn_name_str));
    // `__<PascalName>Args` for the args struct.
    let args_struct = format_ident!("__{}Args", to_pascal_case(&fn_name_str));

    let arg_idents: Vec<_> = func.args.iter().map(|a| &a.ident).collect();
    let arg_types: Vec<_> = func.args.iter().map(|a| &a.ty).collect();

    quote! {
        // Original function — kept callable directly for unit tests
        // and direct use.
        #original_fn

        #[derive(::serde::Deserialize, ::schemars::JsonSchema)]
        #[allow(non_camel_case_types, non_snake_case)]
        struct #args_struct {
            #( #arg_idents: #arg_types, )*
        }

        /// Macro-generated `Tool` wrapper.
        #[derive(Debug, Default, Clone, Copy)]
        pub struct #struct_name;

        impl #struct_name {
            fn cached_schema() -> &'static ::serde_json::Value {
                static SCHEMA: ::std::sync::OnceLock<::serde_json::Value> =
                    ::std::sync::OnceLock::new();
                SCHEMA.get_or_init(|| {
                    let s = ::schemars::schema_for!(#args_struct);
                    ::serde_json::to_value(&s)
                        .expect("schemars produces a valid serde_json::Value")
                })
            }
        }

        // W3.A18: paths route through `::klieo_core::__macro_support`
        // (a `#[doc(hidden)]` re-export layer) so renaming a submodule
        // upstream does not break downstream `#[tool]` invocations.
        #[::async_trait::async_trait]
        impl ::klieo_core::__macro_support::Tool for #struct_name {
            fn name(&self) -> &str {
                #tool_name_str
            }
            fn description(&self) -> &str {
                #description_str
            }
            fn json_schema(&self) -> &::serde_json::Value {
                Self::cached_schema()
            }
            async fn invoke(
                &self,
                args: ::serde_json::Value,
                ctx: ::klieo_core::__macro_support::ToolCtx,
            ) -> ::std::result::Result<
                ::serde_json::Value,
                ::klieo_core::__macro_support::ToolError,
            > {
                let parsed: #args_struct = ::serde_json::from_value(args)
                    .map_err(|e| ::klieo_core::__macro_support::ToolError::InvalidArgs(e.to_string()))?;
                let result = #fn_name(&ctx, #( parsed.#arg_idents ),*).await?;
                ::serde_json::to_value(result)
                    .map_err(|e| ::klieo_core::__macro_support::ToolError::Permanent(e.to_string()))
            }
        }
    }
}

fn to_pascal_case(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    let mut capitalize = true;
    for c in s.chars() {
        if c == '_' {
            capitalize = true;
        } else if capitalize {
            out.extend(c.to_uppercase());
            capitalize = false;
        } else {
            out.push(c);
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn pascal_case_simple() {
        assert_eq!(to_pascal_case("greet"), "Greet");
    }

    #[test]
    fn pascal_case_snake() {
        assert_eq!(to_pascal_case("web_search"), "WebSearch");
    }

    #[test]
    fn pascal_case_already_pascal() {
        assert_eq!(to_pascal_case("Greet"), "Greet");
    }
}