Skip to main content

crabtalk_command_codegen/
lib.rs

1use heck::ToKebabCase;
2use proc_macro::TokenStream;
3use quote::{format_ident, quote};
4use syn::{ItemStruct, LitStr, Token, parse::Parse, parse_macro_input};
5
6struct CommandArgs {
7    kind: String,
8    name: Option<String>,
9    label: Option<String>,
10}
11
12impl Parse for CommandArgs {
13    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
14        let mut kind = None;
15        let mut name = None;
16        let mut label = None;
17
18        while !input.is_empty() {
19            let ident: syn::Ident = input.parse()?;
20            input.parse::<Token![=]>()?;
21            let value: LitStr = input.parse()?;
22
23            match ident.to_string().as_str() {
24                "kind" => kind = Some(value.value()),
25                "name" => name = Some(value.value()),
26                "label" => label = Some(value.value()),
27                other => {
28                    return Err(syn::Error::new(
29                        ident.span(),
30                        format!("unknown attribute: {other}"),
31                    ));
32                }
33            }
34
35            if !input.is_empty() {
36                input.parse::<Token![,]>()?;
37            }
38        }
39
40        let kind = kind.ok_or_else(|| input.error("missing required attribute: kind"))?;
41        Ok(CommandArgs { kind, name, label })
42    }
43}
44
45#[proc_macro_attribute]
46pub fn command(attr: TokenStream, item: TokenStream) -> TokenStream {
47    let args = parse_macro_input!(attr as CommandArgs);
48    let input = parse_macro_input!(item as ItemStruct);
49
50    let struct_name = &input.ident;
51    let name = args
52        .name
53        .unwrap_or_else(|| struct_name.to_string().to_kebab_case());
54    let label = args.label.unwrap_or_else(|| format!("ai.crabtalk.{name}"));
55
56    let command_enum = format_ident!("{}Command", struct_name);
57
58    let start_doc = format!("Install and start the {name} service.");
59    let stop_doc = format!("Stop and uninstall the {name} service.");
60    let run_doc = format!("Run the {name} service directly (used by launchd/systemd).");
61    let logs_doc = format!("View {name} service logs.");
62
63    let run_arm = match args.kind.as_str() {
64        "mcp" => quote! {
65            #command_enum::Run => {
66                crabtalk_command::run_mcp(self).await?
67            }
68        },
69        "client" => quote! {
70            #command_enum::Run => {
71                self.run().await?
72            }
73        },
74        _ => {
75            return syn::Error::new_spanned(struct_name, "kind must be \"mcp\" or \"client\"")
76                .to_compile_error()
77                .into();
78        }
79    };
80
81    let expanded = quote! {
82        #input
83
84        impl crabtalk_command::Service for #struct_name {
85            fn name(&self) -> &str {
86                #name
87            }
88            fn description(&self) -> &str {
89                env!("CARGO_PKG_DESCRIPTION")
90            }
91            fn label(&self) -> &str {
92                #label
93            }
94        }
95
96        #[derive(Debug, clap::Subcommand)]
97        pub enum #command_enum {
98            #[doc = #start_doc]
99            Start,
100            #[doc = #stop_doc]
101            Stop,
102            #[doc = #run_doc]
103            Run,
104            #[doc = #logs_doc]
105            Logs {
106                #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
107                tail_args: Vec<String>,
108            },
109        }
110
111        impl #struct_name {
112            pub async fn exec(
113                &self,
114                action: #command_enum,
115            ) -> crabtalk_command::anyhow::Result<()> {
116                use crabtalk_command::Service as _;
117                match action {
118                    #command_enum::Start => self.start()?,
119                    #command_enum::Stop => self.stop()?,
120                    #run_arm
121                    #command_enum::Logs { tail_args } => {
122                        self.logs(&tail_args)?
123                    }
124                }
125                Ok(())
126            }
127        }
128
129        impl #command_enum {
130            /// Init tracing, build a tokio runtime, and run the command.
131            pub fn start(self, svc: #struct_name) {
132                crabtalk_command::run(move || async move { svc.exec(self).await });
133            }
134        }
135    };
136
137    expanded.into()
138}