aria_engine/
setup.rs

1//! Instance-level Engine setup (in-memory; does not write engine.yml).
2
3use thiserror::Error;
4
5pub const INTL_SITE: &str = "https://ariacompute.com";
6pub const INTL_UPGRADE: &str = "https://github.com/ariacompute";
7pub const CN_SITE: &str = "https://ariacompute.cn";
8pub const CN_UPGRADE: &str = "https://gitee.com/ariacompute";
9
10#[derive(Debug, Error, Clone, PartialEq, Eq)]
11pub enum SetupError {
12    #[error("invalid compute: {0}")]
13    InvalidCompute(String),
14    #[error("{0}")]
15    InvalidKey(String),
16}
17
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct SetupConfig {
20    // --- Local (router registration) ---
21    pub router: String,
22    pub router_api_key: String,
23    // --- OAuth (Aria Compute) ---
24    pub serve_site: String,
25    pub serve_api_key: String,
26    // --- Hub / compute ---
27    pub site_url: String,
28    pub upgrade_url: String,
29    pub compute: String,
30    pub hf_token: String,
31    pub modelscope_api_token: String,
32}
33
34impl Default for SetupConfig {
35    fn default() -> Self {
36        Self {
37            router: String::new(),
38            router_api_key: String::new(),
39            serve_site: String::new(),
40            serve_api_key: String::new(),
41            site_url: String::new(),
42            upgrade_url: String::new(),
43            compute: "auto".into(),
44            hf_token: String::new(),
45            modelscope_api_token: String::new(),
46        }
47    }
48}
49
50/// Partial merge. `None` fields are omitted.
51#[derive(Debug, Clone, Default)]
52pub struct SetupUpdates {
53    // Local
54    pub router: Option<String>,
55    pub router_api_key: Option<String>,
56    // OAuth
57    pub serve_site: Option<String>,
58    pub serve_api_key: Option<String>,
59    // Hub / compute
60    pub site_url: Option<String>,
61    pub upgrade_url: Option<String>,
62    pub compute: Option<String>,
63    pub hf_token: Option<String>,
64    pub modelscope_api_token: Option<String>,
65}
66
67fn gateway_region(url: &str) -> Option<&'static str> {
68    let lower = url.to_ascii_lowercase();
69    if lower.contains("ariacompute.cn") || lower.contains("gitee.com/ariacompute") {
70        Some("cn")
71    } else if lower.contains("ariacompute.com") || lower.contains("github.com/ariacompute") {
72        Some("intl")
73    } else {
74        None
75    }
76}
77
78fn pair_urls(region: &str) -> (&'static str, &'static str) {
79    if region == "cn" {
80        (CN_SITE, CN_UPGRADE)
81    } else {
82        (INTL_SITE, INTL_UPGRADE)
83    }
84}
85
86/// Fill missing site/upgrade URLs from a provided TLD.
87pub fn fill_setup_urls(mut cfg: SetupConfig) -> SetupConfig {
88    let region = gateway_region(&cfg.site_url).or_else(|| gateway_region(&cfg.upgrade_url));
89    let Some(region) = region else {
90        return cfg;
91    };
92    let (site, upgrade) = pair_urls(region);
93    if cfg.site_url.is_empty() {
94        cfg.site_url = site.into();
95    }
96    if cfg.upgrade_url.is_empty() {
97        cfg.upgrade_url = upgrade.into();
98    }
99    cfg
100}
101
102fn validate_router_api_key(key: &str) -> Result<(), SetupError> {
103    let t = key.trim();
104    if t.is_empty() {
105        return Ok(());
106    }
107    if t.starts_with("bfvk-") {
108        return Err(SetupError::InvalidKey(
109            "OAuth key detected (bfvk-); use serve_api_key / [2/2] OAuth (Aria Compute)".into(),
110        ));
111    }
112    Ok(())
113}
114
115fn validate_serve_api_key(key: &str) -> Result<(), SetupError> {
116    let t = key.trim();
117    if t.is_empty() {
118        return Ok(());
119    }
120    if t.starts_with("sk-aria_") {
121        return Err(SetupError::InvalidKey(
122            "Local router key detected (sk-aria_); use router_api_key / [1/2] Local".into(),
123        ));
124    }
125    if !t.starts_with("bfvk-") {
126        return Err(SetupError::InvalidKey(
127            "serve_api_key must start with bfvk-".into(),
128        ));
129    }
130    Ok(())
131}
132
133/// Merge `updates` into `existing`. Validates; does not mutate `existing`.
134pub fn apply_setup(existing: &SetupConfig, updates: &SetupUpdates) -> Result<SetupConfig, SetupError> {
135    let mut out = existing.clone();
136    if let Some(v) = &updates.router {
137        out.router = v.clone();
138    }
139    if let Some(v) = &updates.router_api_key {
140        validate_router_api_key(v)?;
141        out.router_api_key = v.clone();
142    }
143    if let Some(v) = &updates.serve_site {
144        out.serve_site = v.clone();
145    }
146    if let Some(v) = &updates.serve_api_key {
147        validate_serve_api_key(v)?;
148        out.serve_api_key = v.clone();
149    }
150    if let Some(v) = &updates.site_url {
151        out.site_url = v.clone();
152    }
153    if let Some(v) = &updates.upgrade_url {
154        out.upgrade_url = v.clone();
155    }
156    if let Some(v) = &updates.compute {
157        out.compute = v.clone();
158    }
159    if let Some(v) = &updates.hf_token {
160        out.hf_token = v.clone();
161    }
162    if let Some(v) = &updates.modelscope_api_token {
163        out.modelscope_api_token = v.clone();
164    }
165    match out.compute.as_str() {
166        "auto" | "cpu" | "cuda" => {}
167        other => return Err(SetupError::InvalidCompute(other.into())),
168    }
169    validate_router_api_key(&out.router_api_key)?;
170    validate_serve_api_key(&out.serve_api_key)?;
171    Ok(fill_setup_urls(out))
172}
aria_engine/setup.rs

aria_engine/
setup.rs
