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                /// Increase log verbosity (-v = info, -vv = debug, -vvv = trace).
101                #[arg(short, long, action = clap::ArgAction::Count)]
102                verbose: u8,
103            },
104            #[doc = #stop_doc]
105            Stop,
106            #[doc = #run_doc]
107            Run {
108                /// Increase log verbosity (-v = info, -vv = debug, -vvv = trace).
109                #[arg(short, long, action = clap::ArgAction::Count)]
110                verbose: u8,
111            },
112            #[doc = #logs_doc]
113            Logs {
114                #[arg(trailing_var_arg = true, allow_hyphen_values = true)]
115                tail_args: Vec<String>,
116            },
117        }
118
119        impl #struct_name {
120            pub async fn exec(
121                &self,
122                action: #command_enum,
123            ) -> crabtalk_command::anyhow::Result<()> {
124                use crabtalk_command::Service as _;
125                match action {
126                    #command_enum::Start { verbose } => self.start(verbose)?,
127                    #command_enum::Stop => self.stop()?,
128                    #run_arm
129                    #command_enum::Logs { tail_args } => {
130                        self.logs(&tail_args)?
131                    }
132                }
133                Ok(())
134            }
135        }
136
137        impl #command_enum {
138            /// Init tracing, build a tokio runtime, and run the command.
139            pub fn start(self, svc: #struct_name) {
140                let verbose = match &self {
141                    Self::Run { verbose } | Self::Start { verbose } => *verbose,
142                    _ => 0,
143                };
144                crabtalk_command::run(verbose, move || async move { svc.exec(self).await });
145            }
146        }
147    };
148
149    expanded.into()
150}