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