Skip to main content

studio_worker/
cli.rs

1//! Clap CLI definitions, kept out of `main.rs` so they're testable.
2use clap::{Parser, Subcommand};
3
4#[derive(Parser, Debug)]
5#[command(
6    name = "studio-worker",
7    version,
8    about = "Studio worker — pull-based generation agent (image / llm / audio / video)"
9)]
10pub struct Cli {
11    /// Override the path to config.toml.
12    #[arg(long, global = true)]
13    pub config: Option<String>,
14    #[command(subcommand)]
15    pub command: Command,
16}
17
18#[derive(Subcommand, Debug, PartialEq)]
19pub enum Command {
20    /// Start the heartbeat + claim loop.
21    Run,
22    /// Pre-set registration metadata before the next launch.
23    ///
24    /// On a fresh install you don't need this — `run` and `ui`
25    /// auto-register themselves.  Use it explicitly to:
26    ///   * point the worker at a different studio (`--api-base-url`)
27    ///   * clear local registration state after a rejection or
28    ///     between studios (`--reset`)
29    Register {
30        #[arg(long)]
31        api_base_url: Option<String>,
32        #[arg(long)]
33        reset: bool,
34    },
35    /// Print local config + last heartbeat info.
36    Status,
37    /// Install platform-appropriate auto-start service.
38    InstallService,
39    /// Uninstall the auto-start service.
40    UninstallService,
41    /// Set the VRAM threshold (GB) the worker reports.
42    SetThreshold { gb: f32 },
43    /// Print resolved config + relevant paths.
44    Config,
45    /// Check the release feed for a newer version (does not install).
46    CheckUpdate,
47    /// Launch the desktop UI (requires the `ui` cargo feature).
48    Ui,
49}
50
51impl Command {
52    /// Stable kebab-case label for the subcommand.  Used as the
53    /// structured `command` field in the CLI startup breadcrumb so
54    /// operators can filter `journalctl` by which subcommand a
55    /// process is running.  Matches clap's derived subcommand names.
56    pub fn name(&self) -> &'static str {
57        match self {
58            Command::Run => "run",
59            Command::Register { .. } => "register",
60            Command::Status => "status",
61            Command::InstallService => "install-service",
62            Command::UninstallService => "uninstall-service",
63            Command::SetThreshold { .. } => "set-threshold",
64            Command::Config => "config",
65            Command::CheckUpdate => "check-update",
66            Command::Ui => "ui",
67        }
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use clap::Parser;
75
76    #[test]
77    fn parses_run() {
78        let cli = Cli::parse_from(["studio-worker", "run"]);
79        assert!(matches!(cli.command, Command::Run));
80        assert!(cli.config.is_none());
81    }
82
83    #[test]
84    fn parses_run_with_config_override() {
85        let cli = Cli::parse_from(["studio-worker", "--config", "/etc/x.toml", "run"]);
86        assert_eq!(cli.config.as_deref(), Some("/etc/x.toml"));
87        assert!(matches!(cli.command, Command::Run));
88    }
89
90    #[test]
91    fn parses_register_with_overrides() {
92        let cli = Cli::parse_from([
93            "studio-worker",
94            "register",
95            "--api-base-url",
96            "https://example.invalid",
97        ]);
98        match cli.command {
99            Command::Register {
100                api_base_url,
101                reset,
102            } => {
103                assert_eq!(api_base_url.as_deref(), Some("https://example.invalid"));
104                assert!(!reset);
105            }
106            other => panic!("expected register, got {other:?}"),
107        }
108    }
109
110    #[test]
111    fn parses_register_with_reset() {
112        let cli = Cli::parse_from(["studio-worker", "register", "--reset"]);
113        match cli.command {
114            Command::Register {
115                api_base_url,
116                reset,
117            } => {
118                assert!(api_base_url.is_none());
119                assert!(reset);
120            }
121            other => panic!("expected register, got {other:?}"),
122        }
123    }
124
125    #[test]
126    fn parses_bare_register() {
127        let cli = Cli::parse_from(["studio-worker", "register"]);
128        assert!(matches!(
129            cli.command,
130            Command::Register {
131                api_base_url: None,
132                reset: false,
133            }
134        ));
135    }
136
137    #[test]
138    fn parses_set_threshold_with_float() {
139        let cli = Cli::parse_from(["studio-worker", "set-threshold", "12.5"]);
140        match cli.command {
141            Command::SetThreshold { gb } => assert!((gb - 12.5).abs() < 1e-6),
142            other => panic!("expected set-threshold, got {other:?}"),
143        }
144    }
145
146    #[test]
147    fn name_is_stable_kebab_case_for_every_subcommand() {
148        assert_eq!(Command::Run.name(), "run");
149        assert_eq!(
150            Command::Register {
151                api_base_url: None,
152                reset: false
153            }
154            .name(),
155            "register"
156        );
157        assert_eq!(Command::Status.name(), "status");
158        assert_eq!(Command::InstallService.name(), "install-service");
159        assert_eq!(Command::UninstallService.name(), "uninstall-service");
160        assert_eq!(Command::SetThreshold { gb: 1.0 }.name(), "set-threshold");
161        assert_eq!(Command::Config.name(), "config");
162        assert_eq!(Command::CheckUpdate.name(), "check-update");
163        assert_eq!(Command::Ui.name(), "ui");
164    }
165
166    #[test]
167    fn parses_all_simple_subcommands() {
168        let cases = [
169            ("status", Command::Status),
170            ("install-service", Command::InstallService),
171            ("uninstall-service", Command::UninstallService),
172            ("config", Command::Config),
173            ("check-update", Command::CheckUpdate),
174            ("ui", Command::Ui),
175        ];
176        for (name, expected) in cases {
177            let cli = Cli::parse_from(["studio-worker", name]);
178            assert_eq!(cli.command, expected, "parsing `{name}`");
179        }
180    }
181}