klieo-macros 2.2.0

Procedural macros for the klieo agent framework: #[tool] derives a Tool impl from an async fn.
Documentation
//! Parse the `#[tool(...)]` attribute + the decorated function.

use syn::parse::{Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn::{Error, FnArg, Ident, ItemFn, Lit, LitStr, MetaNameValue, ReturnType, Token, Type};

/// Parsed `#[tool(...)]` attribute. Both fields are optional; codegen
/// falls back to the function's first doc-comment line if `description`
/// is absent, and emits a compile error when neither is present.
#[derive(Default)]
pub(crate) struct ToolAttr {
    pub description: Option<LitStr>,
    pub rename: Option<LitStr>,
}

// Manual Debug so test `unwrap_err()` compiles. LitStr doesn't derive Debug,
// but we only need enough for the error path.
impl std::fmt::Debug for ToolAttr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ToolAttr")
            .field("description", &self.description.as_ref().map(|s| s.value()))
            .field("rename", &self.rename.as_ref().map(|s| s.value()))
            .finish()
    }
}

impl Parse for ToolAttr {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let metas: Punctuated<MetaNameValue, Token![,]> = Punctuated::parse_terminated(input)?;
        let mut description: Option<LitStr> = None;
        let mut rename: Option<LitStr> = None;
        for m in metas {
            let key = m
                .path
                .get_ident()
                .ok_or_else(|| Error::new_spanned(&m.path, "expected identifier key"))?
                .clone();
            let value_lit = match m.value {
                syn::Expr::Lit(syn::ExprLit {
                    lit: Lit::Str(s), ..
                }) => s,
                _ => {
                    return Err(Error::new(
                        key.span(),
                        format!("`{key}` must be a string literal"),
                    ));
                }
            };
            if key == "description" {
                description = Some(value_lit);
            } else if key == "rename" {
                rename = Some(value_lit);
            } else {
                return Err(Error::new(
                    key.span(),
                    format!("unknown #[tool] attribute key {key}"),
                ));
            }
        }
        Ok(ToolAttr {
            description,
            rename,
        })
    }
}

/// Extract the first non-empty trimmed line from `#[doc = "..."]`
/// attributes on the function. `///` comments are canonicalised to this
/// form by rustc before macros see them.
pub(crate) fn first_doc_line(attrs: &[syn::Attribute]) -> Option<String> {
    for a in attrs {
        if !a.path().is_ident("doc") {
            continue;
        }
        let nv: syn::MetaNameValue = match a.meta.require_name_value() {
            Ok(nv) => nv.clone(),
            Err(_) => continue,
        };
        if let syn::Expr::Lit(syn::ExprLit {
            lit: Lit::Str(s), ..
        }) = nv.value
        {
            let trimmed = s.value().trim().to_string();
            if !trimmed.is_empty() {
                return Some(trimmed);
            }
        }
    }
    None
}

/// One non-ctx argument extracted from the function signature.
// Fields consumed by Task 3 codegen.
#[allow(dead_code)]
pub(crate) struct ToolArg {
    pub ident: Ident,
    pub ty: Type,
}

/// Parsed function signature.
// Fields consumed by Task 3 codegen.
#[allow(dead_code)]
pub(crate) struct ToolFn {
    pub item_fn: ItemFn,
    pub fn_name: Ident,
    pub args: Vec<ToolArg>,
    pub output_ty: Type,
}

// Manual Debug so test `unwrap_err()` compiles without requiring all
// syn types to be Debug. We only surface the fn name.
impl std::fmt::Debug for ToolFn {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ToolFn")
            .field("fn_name", &self.fn_name.to_string())
            .finish_non_exhaustive()
    }
}

impl ToolFn {
    pub fn parse(item_fn: ItemFn) -> syn::Result<Self> {
        let sig = &item_fn.sig;

        if sig.asyncness.is_none() {
            return Err(Error::new_spanned(
                sig.fn_token,
                "#[tool] requires an async fn",
            ));
        }

        let fn_name = sig.ident.clone();

        // Guard: empty / leading-underscore names produce surprising
        // PascalCase output (e.g. `_hidden` → `Hidden`, dropping the
        // underscore intent). Reject explicitly.
        let name_str = fn_name.to_string();
        if name_str.is_empty() || name_str.starts_with('_') {
            return Err(Error::new_spanned(
                &fn_name,
                "#[tool] does not support empty or leading-underscore function names; use a public-style identifier (e.g. `greet` not `_greet`)",
            ));
        }

        // First arg must be a typed pattern (the ctx). We skip it
        // positionally and don't enforce the exact type — the
        // generated trampoline uses whatever type the user wrote.
        let inputs = &sig.inputs;
        if inputs.is_empty() {
            return Err(Error::new_spanned(
                sig,
                "#[tool] requires at least one arg (ctx: &ToolCtx)",
            ));
        }

        let mut iter = inputs.iter();
        match iter.next().unwrap() {
            FnArg::Receiver(_) => {
                return Err(Error::new_spanned(
                    sig,
                    "#[tool] cannot decorate methods (use a free fn)",
                ));
            }
            FnArg::Typed(_) => {}
        }

        let mut args = Vec::new();
        for fn_arg in iter {
            match fn_arg {
                FnArg::Receiver(_) => {
                    return Err(Error::new_spanned(
                        fn_arg,
                        "#[tool] cannot decorate methods (use a free fn)",
                    ));
                }
                FnArg::Typed(pat_ty) => {
                    let ident = match &*pat_ty.pat {
                        syn::Pat::Ident(pat_ident) => pat_ident.ident.clone(),
                        _ => {
                            return Err(Error::new_spanned(
                                &pat_ty.pat,
                                "#[tool] requires plain `name: Type` arg patterns",
                            ));
                        }
                    };
                    args.push(ToolArg {
                        ident,
                        ty: (*pat_ty.ty).clone(),
                    });
                }
            }
        }

        let output_ty = match &sig.output {
            ReturnType::Default => {
                return Err(Error::new_spanned(
                    &sig.output,
                    "#[tool] requires a return type Result<T, ToolError>",
                ));
            }
            ReturnType::Type(_, ty) => (**ty).clone(),
        };

        Ok(ToolFn {
            item_fn,
            fn_name,
            args,
            output_ty,
        })
    }
}

/// Combined parsed input.
// Fields consumed by Task 3 codegen.
#[allow(dead_code)]
pub(crate) struct ToolDecl {
    pub attr: ToolAttr,
    pub func: ToolFn,
}

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

    #[test]
    fn parses_simple_function() {
        let f: ItemFn = parse2(quote! {
            async fn greet(ctx: &ToolCtx, name: String) -> Result<String, ToolError> {
                Ok(format!("hi {name}"))
            }
        })
        .unwrap();
        let parsed = ToolFn::parse(f).unwrap();
        assert_eq!(parsed.fn_name.to_string(), "greet");
        assert_eq!(parsed.args.len(), 1);
        assert_eq!(parsed.args[0].ident.to_string(), "name");
    }

    #[test]
    fn rejects_non_async() {
        let f: ItemFn = parse2(quote! {
            fn greet(ctx: &ToolCtx, name: String) -> Result<String, ToolError> { unimplemented!() }
        })
        .unwrap();
        let err = ToolFn::parse(f).unwrap_err();
        assert!(err.to_string().contains("async fn"));
    }

    #[test]
    fn rejects_no_return_type() {
        let f: ItemFn = parse2(quote! {
            async fn greet(ctx: &ToolCtx, name: String) {}
        })
        .unwrap();
        let err = ToolFn::parse(f).unwrap_err();
        assert!(err.to_string().contains("return type"));
    }

    #[test]
    fn rejects_no_args() {
        let f: ItemFn = parse2(quote! {
            async fn greet() -> Result<String, ToolError> { unimplemented!() }
        })
        .unwrap();
        let err = ToolFn::parse(f).unwrap_err();
        assert!(err.to_string().contains("at least one arg"));
    }

    #[test]
    fn parses_attr_description() {
        let attr_tokens: proc_macro2::TokenStream = quote! { description = "Greet a person" };
        let attr: ToolAttr = parse2(attr_tokens).unwrap();
        assert_eq!(attr.description.as_ref().unwrap().value(), "Greet a person");
        assert!(attr.rename.is_none());
    }

    #[test]
    fn rejects_unknown_key() {
        let attr_tokens: proc_macro2::TokenStream = quote! { unknown = "x" };
        let err: syn::Error = parse2::<ToolAttr>(attr_tokens).unwrap_err();
        assert!(err.to_string().contains("unknown #[tool] attribute key"));
    }

    #[test]
    fn parses_attr_with_rename_and_no_description() {
        let attr_tokens: proc_macro2::TokenStream = quote! { rename = "explicitName" };
        let attr: ToolAttr = parse2(attr_tokens).unwrap();
        assert!(attr.description.is_none());
        assert_eq!(attr.rename.as_ref().unwrap().value(), "explicitName");
    }

    #[test]
    fn parses_attr_with_no_args() {
        // `#[tool]` alone — both description and rename absent. Parser
        // accepts this; codegen will require a doc comment to fall back on.
        let attr_tokens: proc_macro2::TokenStream = quote! {};
        let attr: ToolAttr = parse2(attr_tokens).unwrap();
        assert!(attr.description.is_none());
        assert!(attr.rename.is_none());
    }

    #[test]
    fn rejects_method_receiver() {
        let f: syn::Result<ItemFn> = parse2(quote! {
            async fn greet(self, name: String) -> Result<String, ToolError> {
                Ok(name)
            }
        });
        // Method-receiver fns might fail at syn parse OR at ToolFn::parse.
        // Both are acceptable — we just need to confirm rejection.
        match f {
            Ok(item_fn) => {
                let err = ToolFn::parse(item_fn).unwrap_err();
                assert!(err.to_string().contains("decorate methods"));
            }
            Err(_) => {
                // syn rejected at parse time; that's also a valid rejection path.
            }
        }
    }

    #[test]
    fn rejects_pattern_arg() {
        let f: ItemFn = parse2(quote! {
            async fn greet(ctx: &ToolCtx, (a, b): (u32, u32)) -> Result<String, ToolError> {
                Ok(String::new())
            }
        })
        .unwrap();
        let err = ToolFn::parse(f).unwrap_err();
        assert!(err.to_string().contains("plain"));
    }

    #[test]
    fn rejects_attr_with_non_string_description() {
        let attr_tokens: proc_macro2::TokenStream = quote! { description = 42 };
        let err = parse2::<ToolAttr>(attr_tokens).unwrap_err();
        assert!(err.to_string().contains("string literal"));
    }

    #[test]
    fn rejects_leading_underscore_fn_name() {
        let f: ItemFn = parse2(quote! {
            async fn _hidden(ctx: &ToolCtx, name: String) -> Result<String, ToolError> {
                Ok(name)
            }
        })
        .unwrap();
        let err = ToolFn::parse(f).unwrap_err();
        assert!(err.to_string().contains("leading-underscore"));
    }
}