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 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
347pub 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
362pub 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
564pub 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 prev_lib = std::env::var("ARIA_FFI_LIB").ok();
756 std::env::remove_var("ARIA_FFI_LIB");
757 let src_dir = tmp.path().join("src");
758 std::fs::create_dir_all(&src_dir).unwrap();
759 let want = if cfg!(windows) {
760 "aria_ffi.dll"
761 } else if cfg!(target_os = "macos") {
762 "libaria_ffi.dylib"
763 } else {
764 "libaria_ffi.so"
765 };
766 std::fs::write(src_dir.join(want), b"dummy-ffi").unwrap();
767 let archive = tmp.path().join("libaria_ffi.tar.gz");
768 {
769 let f = std::fs::File::create(&archive).unwrap();
770 let enc = flate2::write::GzEncoder::new(f, flate2::Compression::default());
771 let mut builder = tar::Builder::new(enc);
772 builder.append_path_with_name(src_dir.join(want), want).unwrap();
773 builder.finish().unwrap();
774 }
775 let dest_dir = tmp.path().join("lib");
776 let got = extract_ffi_archive(&archive, &dest_dir, want).unwrap();
777 assert_eq!(got.file_name().unwrap(), want);
778 assert_eq!(std::fs::read(&got).unwrap(), b"dummy-ffi");
779 let cached = ensure_ffi_lib(None).unwrap();
780 assert_eq!(cached, got);
781 std::env::remove_var("ARIA_COMPUTE_HOME");
782 match prev_lib {
783 Some(v) => std::env::set_var("ARIA_FFI_LIB", v),
784 None => std::env::remove_var("ARIA_FFI_LIB"),
785 }
786 }
787}