use proc_macro::TokenStream;
use proc_macro_crate::{FoundCrate, crate_name};
use proc_macro2::Span;
use quote::{format_ident, quote};
use syn::{FnArg, ImplItem, ItemFn, ItemImpl, LitStr, Pat, ReturnType, parse_macro_input};
fn echo_agent_crate_path() -> syn::Result<syn::Path> {
match crate_name("echo_agent").map_err(|e| syn::Error::new(Span::call_site(), e.to_string()))? {
FoundCrate::Itself => Ok(syn::parse_quote!(::echo_agent)),
FoundCrate::Name(name) => {
let ident = syn::Ident::new(&name, Span::call_site());
Ok(syn::parse_quote!(::#ident))
}
}
}
struct ToolAttrs {
name: String,
description: String,
permissions: Vec<syn::Ident>,
}
impl syn::parse::Parse for ToolAttrs {
fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
let mut name: Option<String> = None;
let mut description: Option<String> = None;
let mut permissions: Vec<syn::Ident> = Vec::new();
while !input.is_empty() {
let ident: syn::Ident = input.parse()?;
if ident == "name" {
let _eq: syn::Token![=] = input.parse()?;
let value: LitStr = input.parse()?;
name = Some(value.value());
} else if ident == "description" {
let _eq: syn::Token![=] = input.parse()?;
let value: LitStr = input.parse()?;
description = Some(value.value());
} else if ident == "permissions" {
let _eq: syn::Token![=] = input.parse()?;
let content;
syn::bracketed!(content in input);
while !content.is_empty() {
let perm: syn::Ident = content.parse()?;
permissions.push(perm);
if !content.is_empty() {
let _comma: syn::Token![,] = content.parse()?;
}
}
} else {
return Err(syn::Error::new_spanned(
ident,
"unknown attribute, expected `name`, `description`, or `permissions`",
));
}
if !input.is_empty() {
let _comma: syn::Token![,] = input.parse()?;
}
}
let name = name.ok_or_else(|| {
syn::Error::new(Span::call_site(), "#[tool] requires `name = \"...\"`")
})?;
let description = description.ok_or_else(|| {
syn::Error::new(
Span::call_site(),
"#[tool] requires `description = \"...\"`",
)
})?;
Ok(ToolAttrs {
name,
description,
permissions,
})
}
}
#[proc_macro_attribute]
pub fn tool(attr: TokenStream, item: TokenStream) -> TokenStream {
let attrs = parse_macro_input!(attr as ToolAttrs);
let input_fn = parse_macro_input!(item as ItemFn);
match tool_impl(attrs, input_fn) {
Ok(ts) => ts.into(),
Err(e) => e.to_compile_error().into(),
}
}
fn tool_impl(attrs: ToolAttrs, func: ItemFn) -> syn::Result<proc_macro2::TokenStream> {
let echo_agent = echo_agent_crate_path()?;
let tool_name = &attrs.name;
let tool_desc = &attrs.description;
let fn_name = &func.sig.ident;
let fn_name_str = fn_name.to_string();
let struct_name = format_ident!("{}Tool", to_pascal_case(&fn_name_str));
let params_name = format_ident!("{}Params", to_pascal_case(&fn_name_str));
if let ReturnType::Default = &func.sig.output {
return Err(syn::Error::new_spanned(
&func.sig,
"#[tool] function must have a return type (e.g., Result<ToolResult>)",
));
}
let (param_fields, param_names) = extract_fn_params(&func)?;
let body = &func.block;
let permissions_override = if attrs.permissions.is_empty() {
quote! {}
} else {
let perms = attrs.permissions.iter().map(|p| {
quote! { #echo_agent::tools::permission::ToolPermission::#p }
});
quote! {
fn permissions(&self) -> Vec<#echo_agent::tools::permission::ToolPermission> {
vec![#(#perms),*]
}
}
};
let expanded = quote! {
#[derive(::serde::Deserialize, ::schemars::JsonSchema)]
pub struct #params_name {
#(#param_fields),*
}
pub struct #struct_name;
impl #echo_agent::tools::Tool for #struct_name {
fn name(&self) -> &str { #tool_name }
fn description(&self) -> &str { #tool_desc }
fn parameters(&self) -> ::serde_json::Value {
let schema = ::schemars::schema_for!(#params_name);
::serde_json::to_value(schema).unwrap_or_default()
}
fn execute(&self, parameters: #echo_agent::tools::ToolParameters) -> ::futures::future::BoxFuture<'_, #echo_agent::error::Result<#echo_agent::tools::ToolResult>> {
Box::pin(async move {
let value = ::serde_json::Value::Object(parameters.into_iter().collect());
let params: #params_name = ::serde_json::from_value(value)
.map_err(|e| #echo_agent::error::ToolError::InvalidParameter {
name: "(deserialization)".to_string(),
message: e.to_string(),
})?;
let #params_name { #(#param_names),* } = params;
#body
})
}
}
impl #struct_name {
#permissions_override
}
};
Ok(expanded)
}
#[proc_macro_attribute]
pub fn callback(_attr: TokenStream, item: TokenStream) -> TokenStream {
let input = parse_macro_input!(item as ItemImpl);
match callback_impl(input) {
Ok(ts) => ts.into(),
Err(e) => e.to_compile_error().into(),
}
}
fn callback_impl(input: ItemImpl) -> syn::Result<proc_macro2::TokenStream> {
let echo_agent = echo_agent_crate_path()?;
let self_ty = &input.self_ty;
let method_impls = impl_block_to_boxfuture_methods(&input, "()")?;
Ok(quote! {
impl #echo_agent::agent::AgentCallback for #self_ty {
#(#method_impls)*
}
})
}
struct NameAttr {
name: String,
}
impl syn::parse::Parse for NameAttr {
fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
let mut name: Option<String> = None;
while !input.is_empty() {
let ident: syn::Ident = input.parse()?;
let _eq: syn::Token![=] = input.parse()?;
let value: LitStr = input.parse()?;
if ident == "name" {
name = Some(value.value());
} else {
return Err(syn::Error::new_spanned(ident, "expected `name`"));
}
if !input.is_empty() {
let _: syn::Token![,] = input.parse()?;
}
}
let name =
name.ok_or_else(|| syn::Error::new(Span::call_site(), "requires `name = \"...\"`"))?;
Ok(NameAttr { name })
}
}
#[proc_macro_attribute]
pub fn guard(attr: TokenStream, item: TokenStream) -> TokenStream {
let attrs = parse_macro_input!(attr as NameAttr);
let input_fn = parse_macro_input!(item as ItemFn);
match guard_impl(attrs, input_fn) {
Ok(ts) => ts.into(),
Err(e) => e.to_compile_error().into(),
}
}
fn guard_impl(attrs: NameAttr, func: ItemFn) -> syn::Result<proc_macro2::TokenStream> {
let echo_agent = echo_agent_crate_path()?;
let guard_name = &attrs.name;
let struct_name = format_ident!("{}Guard", to_pascal_case(&guard_name.replace('-', "_")));
require_return_type(&func)?;
let body = &func.block;
Ok(quote! {
pub struct #struct_name;
impl #echo_agent::guard::Guard for #struct_name {
fn name(&self) -> &str { #guard_name }
fn check<'a>(
&'a self,
content: &'a str,
direction: #echo_agent::guard::GuardDirection,
) -> ::futures::future::BoxFuture<'a, #echo_agent::error::Result<#echo_agent::guard::GuardResult>> {
Box::pin(async move #body)
}
}
})
}
#[proc_macro_attribute]
pub fn handler(_attr: TokenStream, item: TokenStream) -> TokenStream {
let input = parse_macro_input!(item as ItemImpl);
match handler_impl(input) {
Ok(ts) => ts.into(),
Err(e) => e.to_compile_error().into(),
}
}
fn handler_impl(input: ItemImpl) -> syn::Result<proc_macro2::TokenStream> {
let echo_agent = echo_agent_crate_path()?;
let self_ty = &input.self_ty;
let method_impls = extract_boxfuture_methods_with_return(&input)?;
Ok(quote! {
impl #echo_agent::human_loop::HumanLoopHandler for #self_ty {
#(#method_impls)*
}
})
}
#[proc_macro_attribute]
pub fn compressor(_attr: TokenStream, item: TokenStream) -> TokenStream {
let input_fn = parse_macro_input!(item as ItemFn);
match compressor_impl(input_fn) {
Ok(ts) => ts.into(),
Err(e) => e.to_compile_error().into(),
}
}
fn compressor_impl(func: ItemFn) -> syn::Result<proc_macro2::TokenStream> {
let echo_agent = echo_agent_crate_path()?;
let fn_name = &func.sig.ident;
let struct_name = format_ident!("{}Compressor", to_pascal_case(&fn_name.to_string()));
require_return_type(&func)?;
let body = &func.block;
Ok(quote! {
pub struct #struct_name;
impl #echo_agent::compression::ContextCompressor for #struct_name {
fn compress(
&self,
input: #echo_agent::compression::CompressionInput,
) -> ::futures::future::BoxFuture<'_, #echo_agent::error::Result<#echo_agent::compression::CompressionOutput>> {
Box::pin(async move #body)
}
}
})
}
#[proc_macro_attribute]
pub fn permission_policy(_attr: TokenStream, item: TokenStream) -> TokenStream {
let input_fn = parse_macro_input!(item as ItemFn);
match permission_policy_impl(input_fn) {
Ok(ts) => ts.into(),
Err(e) => e.to_compile_error().into(),
}
}
fn permission_policy_impl(func: ItemFn) -> syn::Result<proc_macro2::TokenStream> {
let echo_agent = echo_agent_crate_path()?;
let fn_name = &func.sig.ident;
let struct_name = format_ident!("{}Policy", to_pascal_case(&fn_name.to_string()));
require_return_type(&func)?;
let body = &func.block;
Ok(quote! {
pub struct #struct_name;
impl #echo_agent::tools::permission::PermissionPolicy for #struct_name {
fn check<'a>(
&'a self,
tool_name: &'a str,
permissions: &'a [#echo_agent::tools::permission::ToolPermission],
) -> ::futures::future::BoxFuture<'a, #echo_agent::tools::permission::PermissionDecision> {
Box::pin(async move #body)
}
}
})
}
#[proc_macro_attribute]
pub fn audit_logger(_attr: TokenStream, item: TokenStream) -> TokenStream {
let input = parse_macro_input!(item as ItemImpl);
match audit_logger_impl(input) {
Ok(ts) => ts.into(),
Err(e) => e.to_compile_error().into(),
}
}
fn audit_logger_impl(input: ItemImpl) -> syn::Result<proc_macro2::TokenStream> {
let echo_agent = echo_agent_crate_path()?;
let self_ty = &input.self_ty;
let method_impls = extract_boxfuture_methods_with_return(&input)?;
Ok(quote! {
impl #echo_agent::audit::AuditLogger for #self_ty {
#(#method_impls)*
}
})
}
fn extract_fn_params(
func: &ItemFn,
) -> syn::Result<(Vec<proc_macro2::TokenStream>, Vec<syn::Ident>)> {
let mut param_fields = Vec::new();
let mut param_names = Vec::new();
for arg in func.sig.inputs.iter() {
if let FnArg::Typed(pat_type) = arg {
let pat = &pat_type.pat;
let ty = &pat_type.ty;
let field_name = if let Pat::Ident(pi) = pat.as_ref() {
pi.ident.clone()
} else {
return Err(syn::Error::new_spanned(pat, "expected identifier pattern"));
};
let doc_str = extract_doc_comments(&pat_type.attrs);
let schemars_attr = if let Some(doc) = &doc_str {
quote! { #[schemars(description = #doc)] }
} else {
quote! {}
};
param_fields.push(quote! {
#schemars_attr
pub #field_name: #ty
});
param_names.push(field_name);
}
}
Ok((param_fields, param_names))
}
fn impl_block_to_boxfuture_methods(
input: &ItemImpl,
_default_return: &str,
) -> syn::Result<Vec<proc_macro2::TokenStream>> {
let mut methods = Vec::new();
for item in &input.items {
if let ImplItem::Fn(method) = item {
let name_ident = &method.sig.ident;
let body = &method.block;
let lifetime_params = lifetimed_params(&method.sig.inputs);
methods.push(quote! {
fn #name_ident<'a>(#(#lifetime_params),*) -> ::futures::future::BoxFuture<'a, ()> {
Box::pin(async move #body)
}
});
}
}
Ok(methods)
}
fn extract_boxfuture_methods_with_return(
input: &ItemImpl,
) -> syn::Result<Vec<proc_macro2::TokenStream>> {
let mut methods = Vec::new();
for item in &input.items {
if let ImplItem::Fn(method) = item {
let name_ident = &method.sig.ident;
let body = &method.block;
let lifetime_params = lifetimed_params(&method.sig.inputs);
let ret_ty = match &method.sig.output {
ReturnType::Default => quote! { () },
ReturnType::Type(_, ty) => quote! { #ty },
};
methods.push(quote! {
fn #name_ident<'a>(#(#lifetime_params),*) -> ::futures::future::BoxFuture<'a, #ret_ty> {
Box::pin(async move #body)
}
});
}
}
Ok(methods)
}
fn lifetimed_params(
inputs: &syn::punctuated::Punctuated<FnArg, syn::token::Comma>,
) -> Vec<proc_macro2::TokenStream> {
inputs
.iter()
.map(|arg| match arg {
FnArg::Receiver(_) => quote! { &'a self },
FnArg::Typed(pat_type) => {
let pat = &pat_type.pat;
let ty = add_lifetime_a(&pat_type.ty);
quote! { #pat: #ty }
}
})
.collect()
}
fn add_lifetime_a(ty: &syn::Type) -> proc_macro2::TokenStream {
match ty {
syn::Type::Reference(r) => {
let elem = &r.elem;
let lifetime = r
.lifetime
.as_ref()
.map(|lt| quote! { #lt })
.unwrap_or(quote! { 'a });
if r.mutability.is_some() {
quote! { &#lifetime mut #elem }
} else {
quote! { &#lifetime #elem }
}
}
other => quote! { #other },
}
}
fn require_return_type(func: &ItemFn) -> syn::Result<()> {
if let ReturnType::Default = &func.sig.output {
return Err(syn::Error::new_spanned(
&func.sig,
"function must have an explicit return type",
));
}
Ok(())
}
fn extract_doc_comments(attrs: &[syn::Attribute]) -> Option<String> {
let docs: Vec<String> = attrs
.iter()
.filter_map(|attr| {
if !attr.path().is_ident("doc") {
return None;
}
if let syn::Meta::NameValue(nv) = &attr.meta
&& let syn::Expr::Lit(expr_lit) = &nv.value
&& let syn::Lit::Str(s) = &expr_lit.lit
{
return Some(s.value().trim().to_string());
}
None
})
.collect();
if docs.is_empty() {
None
} else {
Some(docs.join(" "))
}
}
fn to_pascal_case(s: &str) -> String {
s.split('_')
.map(|word| {
let mut chars = word.chars();
match chars.next() {
None => String::new(),
Some(c) => {
let upper: String = c.to_uppercase().collect();
upper + &chars.collect::<String>()
}
}
})
.collect()
}