cortex-derive 1.0.2

Proc macro for cortex-engine — #[derive(Tool)]
Documentation
//! Cortex derive macros — `#[derive(Tool)]`

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use syn::{
    parse_macro_input, Data, DeriveInput, Expr, Field, Fields, Lit, Meta, Type,
};
use syn::punctuated::Punctuated;
use syn::token::Comma;

/// 为 struct 生成 `Tool` trait 实现。
///
/// 自动生成 `name()`、`description()`、`parameters()` 和 `execute()`。
/// `execute()` 将参数反序列化后调用 `__tool_execute_impl(...)`。
#[proc_macro_derive(Tool, attributes(tool))]
pub fn derive_tool(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);

    let fields = match &input.data {
        Data::Struct(data) => match &data.fields {
            Fields::Named(named) => &named.named,
            _ => return err_compile("Tool 只支持具名字段的 struct"),
        },
        _ => return err_compile("Tool 只能 derive 在 struct 上"),
    };

    let struct_name = &input.ident;
    let (tool_name, tool_desc) = parse_tool_attrs(&input);

    let mut field_idents: Vec<TokenStream2> = Vec::new();
    let mut schema_entries = Vec::new();
    let mut deser_code = Vec::new();
    let mut required_fields: Vec<String> = Vec::new();

    for field in fields.iter() {
        let fname = field.ident.as_ref().unwrap();
        let fname_str = fname.to_string();
        let desc = field_doc(field).unwrap_or_else(|| format!("字段 {}", fname_str));
        let (schema_json, deser, required) = type_info(&field.ty, &fname_str, &desc);

        if required {
            required_fields.push(fname_str.clone());
        }
        field_idents.push(quote! { #fname });
        schema_entries.push(quote! {
            #fname_str: #schema_json
        });
        deser_code.push(deser);
    }

    let schema = if required_fields.is_empty() {
        quote! {
            ::serde_json::json!({
                "type": "object",
                "properties": { #(#schema_entries,)* }
            })
        }
    } else {
        quote! {
            ::serde_json::json!({
                "type": "object",
                "properties": { #(#schema_entries,)* },
                "required": [#(#required_fields,)*]
            })
        }
    };

    let expanded = quote! {
        #[::async_trait::async_trait]
        impl ::cortex_engine::tool::Tool for #struct_name {
            fn name(&self) -> &str { #tool_name }
            fn description(&self) -> String { #tool_desc.into() }
            fn parameters(&self) -> ::serde_json::Value { #schema }
            async fn execute(
                &self,
                __args: ::serde_json::Value,
            ) -> Result<::serde_json::Value, ::cortex_engine::tool::ToolError> {
                #(#deser_code)*
                self.__tool_execute_impl(#(#field_idents,)*).await
            }
        }
    };

    expanded.into()
}

fn err_compile(msg: &str) -> TokenStream {
    syn::Error::new(proc_macro2::Span::call_site(), msg)
        .to_compile_error()
        .into()
}

/// 解析 #[tool(name = "...", desc = "...")]
fn parse_tool_attrs(input: &DeriveInput) -> (String, String) {
    let default_name = camel_to_snake(&input.ident.to_string());
    let mut tool_name: Option<String> = None;
    let mut tool_desc: Option<String> = None;

    for attr in &input.attrs {
        if !attr.path().is_ident("tool") { continue; }
        if let Ok(metas) = attr.parse_args_with(
            Punctuated::<Meta, Comma>::parse_terminated,
        ) {
            for meta in &metas {
                if let Meta::NameValue(kv) = meta {
                    let key = kv.path.get_ident().map(|i| i.to_string());
                    let value = match &kv.value {
                        Expr::Lit(lit) => match &lit.lit {
                            Lit::Str(s) => s.value(),
                            _ => continue,
                        },
                        _ => continue,
                    };
                    match key.as_deref() {
                        Some("name") => tool_name = Some(value),
                        Some("desc") => tool_desc = Some(value),
                        _ => {}
                    }
                }
            }
        }
    }

    let name = tool_name.unwrap_or(default_name);
    let desc = tool_desc.unwrap_or_else(|| format!("工具: {}", input.ident));
    (name, desc)
}

fn camel_to_snake(s: &str) -> String {
    let mut r = String::with_capacity(s.len() + 4);
    for (i, c) in s.chars().enumerate() {
        if c.is_uppercase() {
            if i > 0 { r.push('_'); }
            for l in c.to_lowercase() { r.push(l); }
        } else { r.push(c); }
    }
    r
}

/// 提取 struct 字段的 doc comment(支持 `/// text` 和 `#[doc = "text"]`)
fn field_doc(field: &Field) -> Option<String> {
    let docs = get_doc_comments(&field.attrs);
    if docs.is_empty() { None } else { Some(docs) }
}

/// 从属性列表中提取所有 `#[doc = "..."]` 并合并
fn get_doc_comments(attrs: &[syn::Attribute]) -> String {
    let mut parts = Vec::new();
    for attr in attrs {
        if !attr.path().is_ident("doc") { continue; }
        // #[doc = "text"] → 解析赋值语句
        // tokens: `= "text"`
        if let Ok(Meta::NameValue(kv)) = attr.parse_args::<Meta>() {
            if let Expr::Lit(el) = &kv.value {
                if let Lit::Str(s) = &el.lit {
                    let val = s.value().trim().to_string();
                    if !val.is_empty() {
                        parts.push(val);
                    }
                }
            }
        }
    }
    parts.join(" ")
}

/// 获取类型对应的 JSON Schema(含 desc)+ 反序列化代码 + 是否必填
fn type_info(ty: &Type, field_name: &str, desc: &str) -> (TokenStream2, TokenStream2, bool) {
    let ident = syn::Ident::new(field_name, proc_macro2::Span::call_site());
    let key_str = field_name.to_string();
    let type_str = quote!(#ty).to_string();
    let desc = desc.to_string();

    match type_str.as_str() {
        "String" | "std :: string :: String" => (
            quote! { ::serde_json::json!({"type": "string", "description": #desc}) },
            quote! {
                let #ident: String = __args.get(#key_str)
                    .and_then(|v| v.as_str()).unwrap_or("").to_string();
            },
            true,
        ),
        "i64" => (
            quote! { ::serde_json::json!({"type": "integer", "description": #desc}) },
            quote! {
                let #ident: i64 = __args.get(#key_str)
                    .and_then(|v| v.as_i64()).unwrap_or(0);
            },
            true,
        ),
        "i32" => (
            quote! { ::serde_json::json!({"type": "integer", "description": #desc}) },
            quote! {
                let #ident: i32 = __args.get(#key_str)
                    .and_then(|v| v.as_i64()).unwrap_or(0) as i32;
            },
            true,
        ),
        "f64" => (
            quote! { ::serde_json::json!({"type": "number", "description": #desc}) },
            quote! {
                let #ident: f64 = __args.get(#key_str)
                    .and_then(|v| v.as_f64()).unwrap_or(0.0);
            },
            true,
        ),
        "f32" => (
            quote! { ::serde_json::json!({"type": "number", "description": #desc}) },
            quote! {
                let #ident: f32 = __args.get(#key_str)
                    .and_then(|v| v.as_f64()).unwrap_or(0.0) as f32;
            },
            true,
        ),
        "bool" => (
            quote! { ::serde_json::json!({"type": "boolean", "description": #desc}) },
            quote! {
                let #ident: bool = __args.get(#key_str)
                    .and_then(|v| v.as_bool()).unwrap_or(false);
            },
            true,
        ),
        // ---------- Option 包装 ----------
        "Option < String >" | "Option < std :: string :: String >" => (
            quote! { ::serde_json::json!({"type": "string", "description": #desc}) },
            quote! {
                let #ident: Option<String> = __args.get(#key_str)
                    .and_then(|v| v.as_str()).map(String::from);
            },
            false,
        ),
        "Option < i64 >" => (
            quote! { ::serde_json::json!({"type": "integer", "description": #desc}) },
            quote! {
                let #ident: Option<i64> = __args.get(#key_str)
                    .and_then(|v| v.as_i64());
            },
            false,
        ),
        "Option < f64 >" => (
            quote! { ::serde_json::json!({"type": "number", "description": #desc}) },
            quote! {
                let #ident: Option<f64> = __args.get(#key_str)
                    .and_then(|v| v.as_f64());
            },
            false,
        ),
        "Option < bool >" => (
            quote! { ::serde_json::json!({"type": "boolean", "description": #desc}) },
            quote! {
                let #ident: Option<bool> = __args.get(#key_str)
                    .and_then(|v| v.as_bool());
            },
            false,
        ),
        // ---------- Vec 包装 ----------
        "Vec < String >" | "Vec < std :: string :: String >" => (
            quote! { ::serde_json::json!({"type": "array", "items": {"type": "string"}, "description": #desc}) },
            quote! {
                let #ident: Vec<String> = __args.get(#key_str)
                    .and_then(|v| v.as_array())
                    .map(|a| a.iter().filter_map(|v| v.as_str().map(String::from)).collect())
                    .unwrap_or_default();
            },
            true,
        ),
        "Vec < i64 >" => (
            quote! { ::serde_json::json!({"type": "array", "items": {"type": "integer"}, "description": #desc}) },
            quote! {
                let #ident: Vec<i64> = __args.get(#key_str)
                    .and_then(|v| v.as_array())
                    .map(|a| a.iter().filter_map(|v| v.as_i64()).collect())
                    .unwrap_or_default();
            },
            true,
        ),
        "Vec < f64 >" => (
            quote! { ::serde_json::json!({"type": "array", "items": {"type": "number"}, "description": #desc}) },
            quote! {
                let #ident: Vec<f64> = __args.get(#key_str)
                    .and_then(|v| v.as_array())
                    .map(|a| a.iter().filter_map(|v| v.as_f64()).collect())
                    .unwrap_or_default();
            },
            true,
        ),
        // ---------- 通配:serde_json::from_value ----------
        _ => (
            quote! { ::serde_json::json!({"type": "string", "description": #desc}) },
            quote! {
                let #ident = ::serde_json::from_value::<_>(
                    __args.get(#key_str).cloned()
                        .unwrap_or(::serde_json::Value::Null),
                ).map_err(|e| ::cortex_engine::tool::ToolError::new(
                    format!("参数 '{}' 格式错误: {}", #key_str, e)
                ))?;
            },
            true,
        ),
    }
}