Skip to main content

aria_engine/
download.rs

1//! Regional public-hub model auto-download for the Rust SDK.
2//!
3//! Matches `aria-engine download`: `.com` → Hugging Face, `.cn` → ModelScope.
4//! Dashboard zip meta is not used. A Dashboard `sk-` / `bfvk-` token is ignored
5//! for hub auth.
6
7use std::io::{Read, Write};
8use std::path::{Path, PathBuf};
9use std::time::Duration;
10use thiserror::Error;
11
12#[derive(Debug, Error)]
13pub enum DownloadError {
14    #[error("invalid model name: {0}")]
15    InvalidModelName(String),
16    #[error("{0}")]
17    Request(String),
18    #[error("download stream failed: {0}")]
19    Stream(String),
20    #[error("invalid bundle after download: {0}")]
21    InvalidBundle(String),
22    #[error("io error: {0}")]
23    Io(#[from] std::io::Error),
24}
25
26const DEFAULT_SITE: &str = "https://ariacompute.com";
27const DEFAULT_SDK: &str = "v1.0";
28const HUB_REQUIRED: &[&str] = &["config.json", "weight.bin"];
29const HUB_OPTIONAL: &[&str] = &[
30    "tokenizer.json",
31    "tokenizer.model",
32    "tokenizer_config.json",
33    "special_tokens_map.json",
34    "vocab.json",
35    "merges.txt",
36];
37
38fn aria_home() -> Result<PathBuf, DownloadError> {
39    if let Ok(override_home) = std::env::var("ARIA_COMPUTE_HOME") {
40        if !override_home.is_empty() {
41            return Ok(PathBuf::from(override_home));
42        }
43    }
44    let home = if cfg!(windows) {
45        std::env::var("USERPROFILE").map_err(|_| {
46            DownloadError::Io(std::io::Error::new(
47                std::io::ErrorKind::NotFound,
48                "could not resolve home directory",
49            ))
50        })?
51    } else {
52        std::env::var("HOME").map_err(|_| {
53            DownloadError::Io(std::io::Error::new(
54                std::io::ErrorKind::NotFound,
55                "could not resolve home directory",
56            ))
57        })?
58    };
59    Ok(PathBuf::from(home).join(".ariacompute"))
60}
61
62fn models_dir() -> Result<PathBuf, DownloadError> {
63    Ok(aria_home()?.join("models"))
64}
65
66/// Parse a model name such as `gemma-4-e2b-it_q4` into `(slug, quant)`.
67/// Quant follows the `_q4`/`_q8`/`_q326`/`_q3.26` suffix; defaults to `int4`.
68/// Optional codebook-share `_channel` / `_group` (e.g. `*_q326_channel`) is ignored.
69fn parse_bundle_name(model: &str) -> Result<(String, String), DownloadError> {
70    if model.is_empty() || model.contains('/') || model.contains('\\') {
71        return Err(DownloadError::InvalidModelName(model.to_string()));
72    }
73    let (slug, quant) = if let Some(idx) = model.rfind("_q") {
74        let (slug, suffix) = model.split_at(idx);
75        let mut suffix = &suffix[2..];
76        if let Some(core) = suffix.strip_suffix("_channel") {
77            suffix = core;
78        } else if let Some(core) = suffix.strip_suffix("_group") {
79            suffix = core;
80        }
81        let quant = match suffix {
82            "4" => "int4",
83            "8" => "int8",
84            "326" | "3.26" => "int326",
85            other => {
86                return Err(DownloadError::InvalidModelName(format!(
87                    "unknown quant suffix _q{other}"
88                )))
89            }
90        };
91        (slug.to_string(), quant.to_string())
92    } else {
93        (model.to_string(), "int4".to_string())
94    };
95    if slug.is_empty() {
96        return Err(DownloadError::InvalidModelName(model.to_string()));
97    }
98    Ok((slug, quant))
99}
100
101fn preferred_public_hub(site: Option<&str>) -> &'static str {
102    if site
103        .unwrap_or("")
104        .to_ascii_lowercase()
105        .contains("ariacompute.cn")
106    {
107        "modelscope"
108    } else {
109        "huggingface"
110    }
111}
112
113fn hub_bearer(token: &str) -> Option<&str> {
114    let t = token.trim();
115    if t.is_empty() {
116        return None;
117    }
118    let low = t.to_ascii_lowercase();
119    if low.starts_with("sk-") || low.starts_with("bfvk-") {
120        None
121    } else {
122        Some(t)
123    }
124}
125
126fn unquote_yaml(v: &str) -> String {
127    let t = v.trim();
128    let b = t.as_bytes();
129    if b.len() >= 2
130        && ((b[0] == b'"' && *b.last().unwrap() == b'"')
131            || (b[0] == b'\'' && *b.last().unwrap() == b'\''))
132    {
133        t[1..t.len() - 1].to_string()
134    } else {
135        t.to_string()
136    }
137}
138
139fn config_yml_scalar(key: &str) -> Option<String> {
140    let path = aria_home().ok()?.join("config.yml");
141    let raw = std::fs::read_to_string(path).ok()?;
142    for line in raw.lines() {
143        if line.starts_with(' ') || line.starts_with('\t') {
144            continue;
145        }
146        let s = line.trim();
147        if s.is_empty() || s.starts_with('#') {
148            continue;
149        }
150        let Some((k, v)) = s.split_once(':') else {
151            continue;
152        };
153        if k.trim() != key {
154            continue;
155        }
156        let val = unquote_yaml(v);
157        if val.is_empty() {
158            return None;
159        }
160        return Some(val);
161    }
162    None
163}
164
165fn hub_token_field(source: &str) -> &'static str {
166    if source == "modelscope" {
167        "modelscope_api_token"
168    } else {
169        "hf_token"
170    }
171}
172
173fn resolve_hub_token(
174    source: &str,
175    token: &str,
176    hf_token: Option<&str>,
177    modelscope_api_token: Option<&str>,
178) -> Option<String> {
179    let named = if source == "modelscope" {
180        modelscope_api_token.unwrap_or("")
181    } else {
182        hf_token.unwrap_or("")
183    };
184    let from_cfg = config_yml_scalar(hub_token_field(source)).unwrap_or_default();
185    for cand in [named, token, from_cfg.as_str()] {
186        if let Some(b) = hub_bearer(cand) {
187            return Some(b.to_string());
188        }
189    }
190    None
191}
192
193fn hub_path_names(model: &str) -> Vec<String> {
194    let mut names = vec![model.to_string()];
195    let mut lower = model.to_ascii_lowercase();
196    let mut core = model.to_string();
197    for suf in ["_channel", "_group"] {
198        if lower.ends_with(suf) {
199            core = model[..model.len() - suf.len()].to_string();
200            lower = core.to_ascii_lowercase();
201            break;
202        }
203    }
204    let mut stems = vec![core.clone()];
205    if lower.ends_with("_q326") {
206        stems.push(format!("{}q3.26", &core[..core.len() - 5]));
207    } else if lower.ends_with("_q3.26") {
208        stems.push(format!("{}q326", &core[..core.len() - 6]));
209    }
210    for stem in stems {
211        for share in ["", "_channel", "_group"] {
212            let cand = format!("{stem}{share}");
213            if !names.iter().any(|n| n == &cand) {
214                names.push(cand);
215            }
216        }
217    }
218    names
219}
220
221fn hub_file_urls(source: &str, model: &str, file: &str) -> Vec<String> {
222    let mut urls = Vec::new();
223    for name in hub_path_names(model) {
224        if source == "modelscope" {
225            for repo in [format!("AriaCompute/{name}"), "AriaCompute/model".into()] {
226                urls.push(format!(
227                    "https://www.modelscope.cn/models/{repo}/resolve/master/{DEFAULT_SDK}/{name}/{file}"
228                ));
229                urls.push(format!(
230                    "https://modelscope.cn/models/{repo}/resolve/master/{DEFAULT_SDK}/{name}/{file}"
231                ));
232            }
233        } else {
234            for repo in [format!("ariacompute/{name}"), "ariacompute/model".into()] {
235                urls.push(format!(
236                    "https://huggingface.co/{repo}/resolve/main/{DEFAULT_SDK}/{name}/{file}"
237                ));
238            }
239        }
240    }
241    urls
242}
243
244fn is_valid_bundle(dir: &Path) -> bool {
245    let weight = dir.join("weight.bin");
246    let config = dir.join("config.json");
247    if !weight.is_file() || !config.is_file() {
248        return false;
249    }
250    let Ok(raw) = std::fs::read_to_string(&config) else {
251        return false;
252    };
253    let Ok(v) = serde_json::from_str::<serde_json::Value>(&raw) else {
254        return false;
255    };
256    v.get("format")
257        .and_then(|x| x.as_str())
258        .map(|f| f == "aria-quant-bundle")
259        .unwrap_or(false)
260}
261
262fn atomic_replace(src: &Path, dst: &Path) -> std::io::Result<()> {
263    if dst.exists() {
264        std::fs::remove_dir_all(dst)?;
265    }
266    std::fs::rename(src, dst)
267}
268
269fn auth_error(source: &str, code: u16) -> DownloadError {
270    let field = if source == "modelscope" {
271        "modelscope_api_token"
272    } else {
273        "hf_token"
274    };
275    DownloadError::Request(format!(
276        "auth failed HTTP {code}; set {field} via aria-engine auth (do not pass a Dashboard sk-/bfvk- key as the hub token)"
277    ))
278}
279
280fn fetch_url_to_file(url: &str, dest: &Path, token: Option<&str>) -> Result<(), DownloadError> {
281    let agent = ureq::AgentBuilder::new()
282        .timeout(Duration::from_secs(600))
283        .build();
284    let mut req = agent.get(url);
285    if let Some(t) = token {
286        req = req.set("Authorization", &format!("Bearer {t}"));
287    }
288    let resp = match req.call() {
289        Ok(r) => r,
290        Err(ureq::Error::Status(code, _)) if code == 401 || code == 403 => {
291            return Err(DownloadError::Request(format!("HTTP {code}")));
292        }
293        Err(ureq::Error::Status(code, _)) => {
294            return Err(DownloadError::Request(format!("HTTP {code}")));
295        }
296        Err(e) => return Err(DownloadError::Request(e.to_string())),
297    };
298    if let Some(parent) = dest.parent() {
299        std::fs::create_dir_all(parent)?;
300    }
301    let mut reader = resp.into_reader();
302    let mut out = std::fs::File::create(dest)?;
303    let mut buf = [0u8; 1024 * 1024];
304    loop {
305        let n = reader
306            .read(&mut buf)
307            .map_err(|e| DownloadError::Stream(e.to_string()))?;
308        if n == 0 {
309            break;
310        }
311        out.write_all(&buf[..n])?;
312    }
313    Ok(())
314}
315
316fn fetch_hub_file(
317    source: &str,
318    model: &str,
319    file: &str,
320    dest: &Path,
321    token: Option<&str>,
322    required: bool,
323) -> Result<bool, DownloadError> {
324    let mut last: Option<DownloadError> = None;
325    for url in hub_file_urls(source, model, file) {
326        match fetch_url_to_file(&url, dest, token) {
327            Ok(()) => return Ok(true),
328            Err(DownloadError::Request(msg))
329                if msg.contains("HTTP 401") || msg.contains("HTTP 403") =>
330            {
331                let code = if msg.contains("401") { 401 } else { 403 };
332                return Err(auth_error(source, code));
333            }
334            Err(e) => last = Some(e),
335        }
336    }
337    if required {
338        Err(DownloadError::Request(format!(
339            "{source}: missing {file}{}",
340            last.map(|e| format!(": {e}")).unwrap_or_default()
341        )))
342    } else {
343        Ok(false)
344    }
345}
346
347/// Download `model` from the regional public hub into
348/// `~/.ariacompute/models/{model}`, then return that directory.
349///
350/// If a valid bundle already exists at the cache path, the download is skipped.
351/// Hub auth: explicit `hf_token` / `modelscope_api_token`, then generic `token`,
352/// then `~/.ariacompute/config.yml` (same keys as `aria-engine auth`).
353/// Dashboard `sk-` / `bfvk-` keys are not sent to the hub.
354pub fn download_model(
355    model: &str,
356    token: &str,
357    site: Option<&str>,
358) -> Result<PathBuf, DownloadError> {
359    download_model_auth(model, token, site, None, None)
360}
361
362/// Like [`download_model`], with named hub tokens matching `aria-engine auth`.
363pub fn download_model_auth(
364    model: &str,
365    token: &str,
366    site: Option<&str>,
367    hf_token: Option<&str>,
368    modelscope_api_token: Option<&str>,
369) -> Result<PathBuf, DownloadError> {
370    parse_bundle_name(model)?;
371    let site = site.unwrap_or(DEFAULT_SITE);
372    let source = preferred_public_hub(Some(site));
373    let hub_owned = resolve_hub_token(source, token, hf_token, modelscope_api_token);
374    let hub_token = hub_owned.as_deref();
375    let cache = models_dir()?.join(model);
376
377    if cache.exists() && is_valid_bundle(&cache) {
378        return Ok(cache);
379    }
380
381    let staging = models_dir()?.join(format!(".{}.partial", model));
382    if staging.exists() {
383        std::fs::remove_dir_all(&staging)?;
384    }
385    std::fs::create_dir_all(&staging)?;
386    let result = (|| {
387        for file in HUB_REQUIRED {
388            fetch_hub_file(source, model, file, &staging.join(file), hub_token, true)?;
389        }
390        for extra in HUB_OPTIONAL {
391            let _ = fetch_hub_file(source, model, extra, &staging.join(extra), hub_token, false);
392        }
393        if !is_valid_bundle(&staging) {
394            return Err(DownloadError::InvalidBundle(
395                "need weight.bin + aria-quant-bundle config.json".into(),
396            ));
397        }
398        atomic_replace(&staging, &cache)?;
399        Ok(cache.clone())
400    })();
401    if result.is_err() && staging.exists() {
402        let _ = std::fs::remove_dir_all(&staging);
403    }
404    result
405}
406
407const SDK_UA: &str = "aria-engine-sdk/0.1.0";
408
409fn ffi_lib_name() -> &'static str {
410    if cfg!(windows) {
411        "aria_ffi.dll"
412    } else if cfg!(target_os = "macos") {
413        "libaria_ffi.dylib"
414    } else {
415        "libaria_ffi.so"
416    }
417}
418
419fn lib_dir() -> Result<PathBuf, DownloadError> {
420    Ok(aria_home()?.join("lib"))
421}
422
423fn cached_ffi_path() -> Result<Option<PathBuf>, DownloadError> {
424    let p = lib_dir()?.join(ffi_lib_name());
425    Ok(if p.is_file() { Some(p) } else { None })
426}
427
428pub(crate) fn ffi_asset_os(os: &str, arch: &str) -> Result<&'static str, DownloadError> {
429    match (os, arch) {
430        ("linux", "x86_64") => Ok("linux_x86_64"),
431        ("linux", "aarch64") => Ok("linux_arm64"),
432        ("macos", _) => Ok("macos"),
433        ("windows", "x86_64") => Ok("windows_x86_64"),
434        _ => Err(DownloadError::Request(format!(
435            "unsupported platform {os}/{arch} for libaria_ffi"
436        ))),
437    }
438}
439
440fn strip_v(tag: &str) -> &str {
441    let t = tag.trim();
442    t.strip_prefix('v').or_else(|| t.strip_prefix('V')).unwrap_or(t)
443}
444
445fn parse_semver(tag: &str) -> Option<(u64, u64, u64)> {
446    let core = strip_v(tag).split(['-', '+']).next().unwrap_or("");
447    let mut parts = core.split('.');
448    let major = parts.next()?.parse().ok()?;
449    let minor = parts.next().unwrap_or("0").parse().ok()?;
450    let patch = parts.next().unwrap_or("0").parse().ok()?;
451    Some((major, minor, patch))
452}
453
454pub(crate) fn select_latest_stable(releases: &[serde_json::Value]) -> Result<String, DownloadError> {
455    let mut best_tag: Option<&str> = None;
456    let mut best_key = (0u64, 0u64, 0u64);
457    let mut found = false;
458    for rel in releases {
459        if rel.get("draft").and_then(|v| v.as_bool()).unwrap_or(false)
460            || rel.get("prerelease").and_then(|v| v.as_bool()).unwrap_or(false)
461        {
462            continue;
463        }
464        let tag = rel
465            .get("tag_name")
466            .or_else(|| rel.get("tag"))
467            .and_then(|v| v.as_str())
468            .unwrap_or("");
469        if let Some(parsed) = parse_semver(tag) {
470            if !found || parsed > best_key {
471                best_key = parsed;
472                best_tag = Some(tag);
473                found = true;
474            }
475        }
476    }
477    best_tag
478        .map(|t| strip_v(t).to_string())
479        .ok_or_else(|| DownloadError::Request("no stable release found for libaria_ffi".into()))
480}
481
482fn upgrade_org(site: Option<&str>) -> String {
483    if let Some(cfg) = config_yml_scalar("upgrade_url") {
484        return cfg.trim_end_matches('/').to_string();
485    }
486    let from_cfg = config_yml_scalar("site_url");
487    let hint = site
488        .or(from_cfg.as_deref())
489        .unwrap_or(DEFAULT_SITE)
490        .to_ascii_lowercase();
491    if hint.contains("ariacompute.cn") || hint.contains("gitee.com") {
492        "https://gitee.com/ariacompute".into()
493    } else {
494        "https://github.com/ariacompute".into()
495    }
496}
497
498fn releases_api_url(org: &str) -> String {
499    let owner = org.trim_end_matches('/').rsplit('/').next().unwrap_or("ariacompute");
500    if org.to_ascii_lowercase().contains("gitee.com") {
501        format!("https://gitee.com/api/v5/repos/{owner}/engine/releases?per_page=30")
502    } else {
503        format!("https://api.github.com/repos/{owner}/engine/releases?per_page=30")
504    }
505}
506
507fn http_get_bytes(url: &str) -> Result<Vec<u8>, DownloadError> {
508    let agent = ureq::AgentBuilder::new()
509        .timeout(Duration::from_secs(600))
510        .build();
511    let resp = agent
512        .get(url)
513        .set("User-Agent", SDK_UA)
514        .call()
515        .map_err(|e| DownloadError::Request(e.to_string()))?;
516    let mut reader = resp.into_reader();
517    let mut buf = Vec::new();
518    reader
519        .read_to_end(&mut buf)
520        .map_err(|e| DownloadError::Stream(e.to_string()))?;
521    Ok(buf)
522}
523
524pub(crate) fn extract_ffi_archive(
525    archive: &Path,
526    dest_dir: &Path,
527    want: &str,
528) -> Result<PathBuf, DownloadError> {
529    let file = std::fs::File::open(archive)?;
530    let dec = flate2::read::GzDecoder::new(file);
531    let mut ar = tar::Archive::new(dec);
532    for entry in ar.entries()? {
533        let mut entry = entry?;
534        let name = entry.path()?;
535        let base = name
536            .file_name()
537            .and_then(|s| s.to_str())
538            .unwrap_or("");
539        if base != want {
540            continue;
541        }
542        std::fs::create_dir_all(dest_dir)?;
543        let dest = dest_dir.join(want);
544        {
545            let mut out = std::fs::File::create(&dest)?;
546            std::io::copy(&mut entry, &mut out)?;
547        }
548        #[cfg(unix)]
549        {
550            use std::os::unix::fs::PermissionsExt;
551            let mut perms = std::fs::metadata(&dest)?.permissions();
552            perms.set_mode(0o755);
553            std::fs::set_permissions(&dest, perms)?;
554        }
555        return Ok(dest);
556    }
557    Err(DownloadError::Request(format!(
558        "{} not found in {}",
559        want,
560        archive.display()
561    )))
562}
563
564/// Return a path to libaria_ffi, downloading the latest stable Release if needed.
565pub fn ensure_ffi_lib(site: Option<&str>) -> Result<PathBuf, DownloadError> {
566    if let Ok(env) = std::env::var("ARIA_FFI_LIB") {
567        let p = PathBuf::from(&env);
568        if p.is_file() {
569            return Ok(p);
570        }
571    }
572    if let Some(cached) = cached_ffi_path()? {
573        return Ok(cached);
574    }
575
576    let org = upgrade_org(site);
577    let raw = http_get_bytes(&releases_api_url(&org))?;
578    let releases: Vec<serde_json::Value> = serde_json::from_slice(&raw)
579        .map_err(|e| DownloadError::Request(format!("invalid releases JSON from {org}: {e}")))?;
580    let ver = select_latest_stable(&releases)?;
581    let asset_os = ffi_asset_os(std::env::consts::OS, std::env::consts::ARCH)?;
582    let asset_name = format!("libaria_ffi_{ver}_{asset_os}.tar.gz");
583    let mut url = None;
584    for rel in &releases {
585        let tag = rel
586            .get("tag_name")
587            .or_else(|| rel.get("tag"))
588            .and_then(|v| v.as_str())
589            .unwrap_or("");
590        if strip_v(tag) != ver {
591            continue;
592        }
593        if let Some(assets) = rel.get("assets").and_then(|v| v.as_array()) {
594            for asset in assets {
595                if asset.get("name").and_then(|v| v.as_str()) == Some(asset_name.as_str()) {
596                    url = asset
597                        .get("browser_download_url")
598                        .or_else(|| asset.get("direct_asset_url"))
599                        .and_then(|v| v.as_str())
600                        .map(|s| s.to_string());
601                    break;
602                }
603            }
604        }
605        if url.is_some() {
606            break;
607        }
608    }
609    let url = url.ok_or_else(|| DownloadError::Request(format!("release asset not found: {asset_name}")))?;
610
611    let staging = aria_home()?.join("tmp").join(format!("ffi-{ver}"));
612    if staging.exists() {
613        std::fs::remove_dir_all(&staging)?;
614    }
615    std::fs::create_dir_all(&staging)?;
616    let archive = staging.join(&asset_name);
617    let result = (|| {
618        let bytes = http_get_bytes(&url)?;
619        if let Some(parent) = archive.parent() {
620            std::fs::create_dir_all(parent)?;
621        }
622        std::fs::write(&archive, bytes)?;
623        let dir = lib_dir()?;
624        extract_ffi_archive(&archive, &dir, ffi_lib_name())
625    })();
626    let _ = std::fs::remove_dir_all(&staging);
627    result
628}
629
630#[cfg(test)]
631pub(crate) static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
632
633#[cfg(test)]
634mod tests {
635    use super::*;
636    use super::ENV_LOCK;
637
638    #[test]
639    fn preferred_hub_follows_site_tld() {
640        assert_eq!(
641            preferred_public_hub(Some("https://ariacompute.com")),
642            "huggingface"
643        );
644        assert_eq!(
645            preferred_public_hub(Some("https://ariacompute.cn")),
646            "modelscope"
647        );
648        assert_eq!(preferred_public_hub(None), "huggingface");
649    }
650
651    #[test]
652    fn dashboard_token_not_sent_to_hub() {
653        assert!(hub_bearer("sk-bf-95076ed1-8c1a-4efa-b33c-f52c1d7f9f24").is_none());
654        assert!(hub_bearer("bfvk-test").is_none());
655        assert_eq!(hub_bearer("hf_abc"), Some("hf_abc"));
656    }
657
658    #[test]
659    fn hub_urls_follow_upload_layout() {
660        let hf = hub_file_urls("huggingface", "gemma-4-e2b-it_q4", "config.json");
661        assert!(hf.iter().any(|u| u.contains(
662            "/ariacompute/gemma-4-e2b-it_q4/resolve/main/v1.0/gemma-4-e2b-it_q4/config.json"
663        )));
664        let ms = hub_file_urls("modelscope", "gemma-4-e2b-it_q4", "weight.bin");
665        assert!(ms
666            .iter()
667            .any(|u| u.contains("/v1.0/gemma-4-e2b-it_q4/weight.bin")));
668        assert!(hf.iter().chain(ms.iter()).all(|u| !u.contains("/api/dashboard/")));
669    }
670
671    #[test]
672    fn parse_channel_suffix() {
673        let (slug, quant) = parse_bundle_name("foo_q326_channel").unwrap();
674        assert_eq!(slug, "foo");
675        assert_eq!(quant, "int326");
676    }
677
678    #[test]
679    fn cached_bundle_skips_download() {
680        let _guard = ENV_LOCK.lock().unwrap();
681        let tmp = tempfile::tempdir().unwrap();
682        std::env::set_var("ARIA_COMPUTE_HOME", tmp.path());
683        let cache = tmp.path().join("models").join("foo_q4");
684        std::fs::create_dir_all(&cache).unwrap();
685        std::fs::write(cache.join("weight.bin"), b"x").unwrap();
686        std::fs::write(
687            cache.join("config.json"),
688            br#"{"format":"aria-quant-bundle"}"#,
689        )
690        .unwrap();
691        let got = download_model("foo_q4", "", None).unwrap();
692        assert_eq!(got, cache);
693        std::env::remove_var("ARIA_COMPUTE_HOME");
694    }
695
696    #[test]
697    fn resolve_named_and_config_yml() {
698        let _guard = ENV_LOCK.lock().unwrap();
699        let tmp = tempfile::tempdir().unwrap();
700        std::env::set_var("ARIA_COMPUTE_HOME", tmp.path());
701        std::fs::write(
702            tmp.path().join("config.yml"),
703            "hf_token: hf_from_yml\nmodelscope_api_token: \"ms_from_yml\"\n",
704        )
705        .unwrap();
706        assert_eq!(
707            resolve_hub_token("huggingface", "hf_generic", Some("hf_named"), None).as_deref(),
708            Some("hf_named")
709        );
710        assert_eq!(
711            resolve_hub_token("modelscope", "", None, Some("ms_named")).as_deref(),
712            Some("ms_named")
713        );
714        assert_eq!(
715            resolve_hub_token("huggingface", "", None, None).as_deref(),
716            Some("hf_from_yml")
717        );
718        assert_eq!(
719            resolve_hub_token("modelscope", "", None, None).as_deref(),
720            Some("ms_from_yml")
721        );
722        assert_eq!(
723            resolve_hub_token("huggingface", "sk-bf-not-hub", None, None).as_deref(),
724            Some("hf_from_yml")
725        );
726        std::env::remove_var("ARIA_COMPUTE_HOME");
727    }
728
729    #[test]
730    fn ffi_asset_os_matches_upgrade() {
731        assert_eq!(ffi_asset_os("linux", "x86_64").unwrap(), "linux_x86_64");
732        assert_eq!(ffi_asset_os("linux", "aarch64").unwrap(), "linux_arm64");
733        assert_eq!(ffi_asset_os("macos", "aarch64").unwrap(), "macos");
734        assert_eq!(ffi_asset_os("windows", "x86_64").unwrap(), "windows_x86_64");
735        assert!(ffi_asset_os("linux", "powerpc64").is_err());
736    }
737
738    #[test]
739    fn select_latest_stable_skips_draft_and_prerelease() {
740        let releases = serde_json::json!([
741            {"tag_name": "v0.7.1", "draft": false, "prerelease": false},
742            {"tag_name": "v0.8.0-rc1", "draft": false, "prerelease": true},
743            {"tag_name": "v0.7.2", "draft": false, "prerelease": false},
744            {"tag_name": "v0.9.0", "draft": true, "prerelease": false}
745        ]);
746        let arr = releases.as_array().unwrap();
747        assert_eq!(select_latest_stable(arr).unwrap(), "0.7.2");
748    }
749
750    #[test]
751    fn extract_ffi_and_cached_skip() {
752        let _guard = ENV_LOCK.lock().unwrap();
753        let tmp = tempfile::tempdir().unwrap();
754        std::env::set_var("ARIA_COMPUTE_HOME", tmp.path());
755        let src_dir = tmp.path().join("src");
756        std::fs::create_dir_all(&src_dir).unwrap();
757        let want = if cfg!(windows) {
758            "aria_ffi.dll"
759        } else if cfg!(target_os = "macos") {
760            "libaria_ffi.dylib"
761        } else {
762            "libaria_ffi.so"
763        };
764        std::fs::write(src_dir.join(want), b"dummy-ffi").unwrap();
765        let archive = tmp.path().join("libaria_ffi.tar.gz");
766        {
767            let f = std::fs::File::create(&archive).unwrap();
768            let enc = flate2::write::GzEncoder::new(f, flate2::Compression::default());
769            let mut builder = tar::Builder::new(enc);
770            builder.append_path_with_name(src_dir.join(want), want).unwrap();
771            builder.finish().unwrap();
772        }
773        let dest_dir = tmp.path().join("lib");
774        let got = extract_ffi_archive(&archive, &dest_dir, want).unwrap();
775        assert_eq!(got.file_name().unwrap(), want);
776        assert_eq!(std::fs::read(&got).unwrap(), b"dummy-ffi");
777        let cached = ensure_ffi_lib(None).unwrap();
778        assert_eq!(cached, got);
779        std::env::remove_var("ARIA_COMPUTE_HOME");
780    }
781}