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    let cli_name = format!("crabtalk-{name}");
58
59    let start_doc = format!("Install and start the {name} service.");
60    let stop_doc = format!("Stop and uninstall the {name} service.");
61    let run_doc = format!("Run the {name} service directly (used by launchd/systemd).");
62    let logs_doc = format!("View {name} service logs.");
63
64    let run_arm = match args.kind.as_str() {
65        "mcp" => quote! {
66            #command_enum::Run => {
67                crabtalk_command::run_mcp(self).await?
68            }
69        },
70        "client" => quote! {
71            #command_enum::Run => {
72                self.run().await?
73            }
74        },
75        _ => {
76            return syn::Error::new_spanned(struct_name, "kind must be \"mcp\" or \"client\"")
77                .to_compile_error()
78                .into();
79        }
80    };
81
82    let expanded = quote! {
83        #input
84
85        impl crabtalk_command::Service for #struct_name {
86            fn name(&self) -> &str {
87                #name
88            }
89            fn description(&self) -> &str {
90                env!("CARGO_PKG_DESCRIPTION")
91            }
92            fn label(&self) -> &str {
93                #label
94            }
95        }
96
97        /// CLI wrapper with global `--verbose` flag.
98        #[derive(Debug, clap::Parser)]
99        #[command(name = #cli_name)]
100        pub struct CrabtalkCli {
101            /// Increase log verbosity (-v = info, -vv = debug, -vvv = trace).
102            #[arg(short, long, action = clap::ArgAction::Count, global = true)]
103            pub verbose: u8,
104            #[command(subcommand)]
105            pub action: #command_enum,
106        }
107
108        #[derive(Debug, clap::Subcommand)]
109        pub enum #command_enum {
110            #[doc = #start_doc]
111            Start {
112                /// Re-install even if already installed.
113                #[arg(short, long)]
114                force: bool,
115            },
116            #[doc = #stop_doc]
117            Stop,
118            #[doc = #run_doc]
119            Run,
120            #[doc = #logs_doc]
121            Logs {
122                #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
123                tail_args: Vec<String>,
124            },
125        }
126
127        impl #struct_name {
128            pub async fn exec(
129                &self,
130                action: #command_enum,
131            ) -> crabtalk_command::anyhow::Result<()> {
132                use crabtalk_command::Service as _;
133                match action {
134                    #command_enum::Start { force } => self.start(force)?,
135                    #command_enum::Stop => self.stop()?,
136                    #run_arm
137                    #command_enum::Logs { tail_args } => {
138                        self.logs(&tail_args)?
139                    }
140                }
141                Ok(())
142            }
143        }
144
145        impl CrabtalkCli {
146            /// Init tracing, build a tokio runtime, and run the command.
147            pub fn start(self, svc: #struct_name) {
148                crabtalk_command::run(self.verbose, move || async move {
149                    svc.exec(self.action).await
150                });
151            }
152        }
153    };
154
155    expanded.into()
156}