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    /// Turnkey finish line: install + start the auto-start service and
38    /// print the studio-approval + local-API guidance.  Run this once
39    /// after installing.
40    Setup,
41    /// Install platform-appropriate auto-start service.
42    InstallService,
43    /// Uninstall the auto-start service.
44    UninstallService,
45    /// Set the VRAM threshold (GB) the worker reports.
46    SetThreshold { gb: f32 },
47    /// Print resolved config + relevant paths.
48    Config,
49    /// Check the release feed for a newer version (does not install).
50    CheckUpdate,
51    /// Launch the desktop UI (requires the `ui` cargo feature).
52    Ui,
53}
54
55impl Command {
56    /// Stable kebab-case label for the subcommand.  Used as the
57    /// structured `command` field in the CLI startup breadcrumb so
58    /// operators can filter `journalctl` by which subcommand a
59    /// process is running.  Matches clap's derived subcommand names.
60    pub fn name(&self) -> &'static str {
61        match self {
62            Command::Run => "run",
63            Command::Setup => "setup",
64            Command::Register { .. } => "register",
65            Command::Status => "status",
66            Command::InstallService => "install-service",
67            Command::UninstallService => "uninstall-service",
68            Command::SetThreshold { .. } => "set-threshold",
69            Command::Config => "config",
70            Command::CheckUpdate => "check-update",
71            Command::Ui => "ui",
72        }
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79    use clap::Parser;
80
81    #[test]
82    fn parses_run() {
83        let cli = Cli::parse_from(["studio-worker", "run"]);
84        assert!(matches!(cli.command, Command::Run));
85        assert!(cli.config.is_none());
86    }
87
88    #[test]
89    fn parses_run_with_config_override() {
90        let cli = Cli::parse_from(["studio-worker", "--config", "/etc/x.toml", "run"]);
91        assert_eq!(cli.config.as_deref(), Some("/etc/x.toml"));
92        assert!(matches!(cli.command, Command::Run));
93    }
94
95    #[test]
96    fn parses_register_with_overrides() {
97        let cli = Cli::parse_from([
98            "studio-worker",
99            "register",
100            "--api-base-url",
101            "https://example.invalid",
102        ]);
103        match cli.command {
104            Command::Register {
105                api_base_url,
106                reset,
107            } => {
108                assert_eq!(api_base_url.as_deref(), Some("https://example.invalid"));
109                assert!(!reset);
110            }
111            other => panic!("expected register, got {other:?}"),
112        }
113    }
114
115    #[test]
116    fn parses_register_with_reset() {
117        let cli = Cli::parse_from(["studio-worker", "register", "--reset"]);
118        match cli.command {
119            Command::Register {
120                api_base_url,
121                reset,
122            } => {
123                assert!(api_base_url.is_none());
124                assert!(reset);
125            }
126            other => panic!("expected register, got {other:?}"),
127        }
128    }
129
130    #[test]
131    fn parses_bare_register() {
132        let cli = Cli::parse_from(["studio-worker", "register"]);
133        assert!(matches!(
134            cli.command,
135            Command::Register {
136                api_base_url: None,
137                reset: false,
138            }
139        ));
140    }
141
142    #[test]
143    fn parses_set_threshold_with_float() {
144        let cli = Cli::parse_from(["studio-worker", "set-threshold", "12.5"]);
145        match cli.command {
146            Command::SetThreshold { gb } => assert!((gb - 12.5).abs() < 1e-6),
147            other => panic!("expected set-threshold, got {other:?}"),
148        }
149    }
150
151    #[test]
152    fn name_is_stable_kebab_case_for_every_subcommand() {
153        assert_eq!(Command::Run.name(), "run");
154        assert_eq!(
155            Command::Register {
156                api_base_url: None,
157                reset: false
158            }
159            .name(),
160            "register"
161        );
162        assert_eq!(Command::Status.name(), "status");
163        assert_eq!(Command::Setup.name(), "setup");
164        assert_eq!(Command::InstallService.name(), "install-service");
165        assert_eq!(Command::UninstallService.name(), "uninstall-service");
166        assert_eq!(Command::SetThreshold { gb: 1.0 }.name(), "set-threshold");
167        assert_eq!(Command::Config.name(), "config");
168        assert_eq!(Command::CheckUpdate.name(), "check-update");
169        assert_eq!(Command::Ui.name(), "ui");
170    }
171
172    #[test]
173    fn parses_all_simple_subcommands() {
174        let cases = [
175            ("status", Command::Status),
176            ("setup", Command::Setup),
177            ("install-service", Command::InstallService),
178            ("uninstall-service", Command::UninstallService),
179            ("config", Command::Config),
180            ("check-update", Command::CheckUpdate),
181            ("ui", Command::Ui),
182        ];
183        for (name, expected) in cases {
184            let cli = Cli::parse_from(["studio-worker", name]);
185            assert_eq!(cli.command, expected, "parsing `{name}`");
186        }
187    }
188}