klieo-macros 3.0.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();

    let effectful_impl = if attr.effectful {
        quote! { fn is_effectful(&self) -> bool { true } }
    } else {
        quote! {}
    };
    let redacts_audit_impl = if attr.redacts_audit {
        quote! { fn redacts_audit(&self) -> bool { true } }
    } else {
        quote! {}
    };

    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")
                })
            }
        }

        // 0.16 P5: paths route through `::klieo::__private::klieo_core::__macro_support`
        // so downstream crates need only `klieo` as a direct dep —
        // `klieo-core` no longer needs to appear in the user's
        // Cargo.toml.
        #[::async_trait::async_trait]
        impl ::klieo::__private::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()
            }
            #effectful_impl
            #redacts_audit_impl
            async fn invoke(
                &self,
                args: ::serde_json::Value,
                ctx: ::klieo::__private::klieo_core::__macro_support::ToolCtx,
            ) -> ::std::result::Result<
                ::serde_json::Value,
                ::klieo::__private::klieo_core::__macro_support::ToolError,
            > {
                let parsed: #args_struct = ::serde_json::from_value(args)
                    .map_err(|e| ::klieo::__private::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::__private::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::*;
    use crate::parse::{ToolAttr, ToolDecl, ToolFn};
    use syn::{parse2, ItemFn};

    fn sample_tool_fn() -> ToolFn {
        let f: ItemFn = parse2(quote! {
            /// pays out
            async fn pay(ctx: &ToolCtx, amount: u64) -> Result<u64, ToolError> {
                unimplemented!()
            }
        })
        .unwrap();
        ToolFn::parse(f).unwrap()
    }

    #[test]
    fn effectful_flag_emits_is_effectful_override() {
        let attr = ToolAttr {
            effectful: true,
            ..Default::default()
        };
        let out = expand(ToolDecl { attr, func: sample_tool_fn() }).to_string();
        assert!(out.contains("fn is_effectful"));
        assert!(!out.contains("fn redacts_audit"), "redacts_audit was not requested");
    }

    #[test]
    fn no_flags_emit_no_overrides() {
        let out = expand(ToolDecl {
            attr: ToolAttr::default(),
            func: sample_tool_fn(),
        })
        .to_string();
        assert!(!out.contains("fn is_effectful"));
        assert!(!out.contains("fn redacts_audit"));
    }

    #[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");
    }
}