studio-worker 0.1.2

Pull-based image-generation worker for the minis.gg studio.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

//! Clap CLI definitions, kept out of `main.rs` so they're testable.
use clap::{Parser, Subcommand};

#[derive(Parser, Debug)]
#[command(
    name = "studio-worker",
    version,
    about = "Studio worker — pull-based generation agent (image / llm / audio / video)"
)]
pub struct Cli {
    /// Override the path to config.toml.
    #[arg(long, global = true)]
    pub config: Option<String>,
    #[command(subcommand)]
    pub command: Command,
}

#[derive(Subcommand, Debug, PartialEq)]
pub enum Command {
    /// Start the heartbeat + claim loop.
    Run,
    /// Register the worker against the API (idempotent).
    Register {
        #[arg(long)]
        bootstrap_token: Option<String>,
        #[arg(long)]
        api_base_url: Option<String>,
    },
    /// Print local config + last heartbeat info.
    Status,
    /// Install platform-appropriate auto-start service.
    InstallService,
    /// Uninstall the auto-start service.
    UninstallService,
    /// Enable auto-claim.
    Enable,
    /// Disable auto-claim.
    Disable,
    /// Set the VRAM threshold (GB) the worker reports.
    SetThreshold { gb: f32 },
    /// Print resolved config + relevant paths.
    Config,
    /// Check the release feed for a newer version (does not install).
    CheckUpdate,
    /// Launch the desktop UI (requires the `ui` cargo feature).
    Ui,
}

#[cfg(test)]
mod tests {
    use super::*;
    use clap::Parser;

    #[test]
    fn parses_run() {
        let cli = Cli::parse_from(["studio-worker", "run"]);
        assert!(matches!(cli.command, Command::Run));
        assert!(cli.config.is_none());
    }

    #[test]
    fn parses_run_with_config_override() {
        let cli = Cli::parse_from(["studio-worker", "--config", "/etc/x.toml", "run"]);
        assert_eq!(cli.config.as_deref(), Some("/etc/x.toml"));
        assert!(matches!(cli.command, Command::Run));
    }

    #[test]
    fn parses_register_with_overrides() {
        let cli = Cli::parse_from([
            "studio-worker",
            "register",
            "--bootstrap-token",
            "secret",
            "--api-base-url",
            "https://example.invalid",
        ]);
        match cli.command {
            Command::Register {
                bootstrap_token,
                api_base_url,
            } => {
                assert_eq!(bootstrap_token.as_deref(), Some("secret"));
                assert_eq!(api_base_url.as_deref(), Some("https://example.invalid"));
            }
            other => panic!("expected register, got {other:?}"),
        }
    }

    #[test]
    fn parses_set_threshold_with_float() {
        let cli = Cli::parse_from(["studio-worker", "set-threshold", "12.5"]);
        match cli.command {
            Command::SetThreshold { gb } => assert!((gb - 12.5).abs() < 1e-6),
            other => panic!("expected set-threshold, got {other:?}"),
        }
    }

    #[test]
    fn parses_all_simple_subcommands() {
        let cases = [
            ("status", Command::Status),
            ("install-service", Command::InstallService),
            ("uninstall-service", Command::UninstallService),
            ("enable", Command::Enable),
            ("disable", Command::Disable),
            ("config", Command::Config),
            ("check-update", Command::CheckUpdate),
            ("ui", Command::Ui),
        ];
        for (name, expected) in cases {
            let cli = Cli::parse_from(["studio-worker", name]);
            assert_eq!(cli.command, expected, "parsing `{name}`");
        }
    }
}