aria_engine/
download.rs

1//! Dashboard-only model auto-download for the Rust SDK.
2//!
3//! Mirrors the `dashboard` branch of `openai/src/download.rs`: resolve the
4//! `slug`/`quant` from the model name, request the meta URL with a bearer
5//! token, stream the zip, validate the zip magic, extract (flattening a single
6//! top-level subdir), and verify the resulting bundle.
7
8use std::io::{Read, Write};
9use std::path::{Path, PathBuf};
10use thiserror::Error;
11
12#[derive(Debug, Error)]
13pub enum DownloadError {
14    #[error("invalid model name: {0}")]
15    InvalidModelName(String),
16    #[error("dashboard request failed: {0}")]
17    Request(String),
18    #[error("download stream failed: {0}")]
19    Stream(String),
20    #[error("invalid zip archive: {0}")]
21    BadZip(String),
22    #[error("extraction failed: {0}")]
23    Extract(String),
24    #[error("invalid bundle after download: {0}")]
25    InvalidBundle(String),
26    #[error("io error: {0}")]
27    Io(#[from] std::io::Error),
28}
29
30const DEFAULT_SITE: &str = "https://ariacompute.com";
31
32fn aria_home() -> Result<PathBuf, DownloadError> {
33    if let Ok(override_home) = std::env::var("ARIA_COMPUTE_HOME") {
34        if !override_home.is_empty() {
35            return Ok(PathBuf::from(override_home));
36        }
37    }
38    let home = if cfg!(windows) {
39        std::env::var("USERPROFILE").map_err(|_| {
40            DownloadError::Io(std::io::Error::new(
41                std::io::ErrorKind::NotFound,
42                "could not resolve home directory",
43            ))
44        })?
45    } else {
46        std::env::var("HOME").map_err(|_| {
47            DownloadError::Io(std::io::Error::new(
48                std::io::ErrorKind::NotFound,
49                "could not resolve home directory",
50            ))
51        })?
52    };
53    Ok(PathBuf::from(home).join(".ariacompute"))
54}
55
56fn models_dir() -> Result<PathBuf, DownloadError> {
57    Ok(aria_home()?.join("models"))
58}
59
60/// Parse a model name such as `gemma-4-e2b-it_q4` into `(slug, quant)`.
61/// Quant follows the `_q4`/`_q8`/`_q326`/`_q3.26` suffix; defaults to `int4`.
62fn parse_bundle_name(model: &str) -> Result<(String, String), DownloadError> {
63    if model.is_empty() || model.contains('/') || model.contains('\\') {
64        return Err(DownloadError::InvalidModelName(model.to_string()));
65    }
66    let (slug, quant) = if let Some(idx) = model.rfind("_q") {
67        let (slug, suffix) = model.split_at(idx);
68        let suffix = &suffix[2..];
69        let quant = match suffix {
70            "4" => "int4",
71            "8" => "int8",
72            "326" | "3.26" => "int326",
73            other => {
74                return Err(DownloadError::InvalidModelName(format!(
75                    "unknown quant suffix _q{other}"
76                )))
77            }
78        };
79        (slug.to_string(), quant.to_string())
80    } else {
81        (model.to_string(), "int4".to_string())
82    };
83    if slug.is_empty() {
84        return Err(DownloadError::InvalidModelName(model.to_string()));
85    }
86    Ok((slug, quant))
87}
88
89fn url_encode(s: &str) -> String {
90    let mut out = String::with_capacity(s.len());
91    for b in s.bytes() {
92        match b {
93            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
94                out.push(b as char)
95            }
96            _ => out.push_str(&format!("%{:02X}", b)),
97        }
98    }
99    out
100}
101
102#[derive(serde::Deserialize)]
103struct DashboardMeta {
104    url: String,
105}
106
107fn meta_url(site: &str, slug: &str, quant: &str) -> String {
108    let sdk = "v1.0";
109    format!(
110        "{}/api/dashboard/models/{}/download?quant={}&sdk={}&format=json",
111        site.trim_end_matches('/'),
112        url_encode(slug),
113        url_encode(quant),
114        url_encode(sdk),
115    )
116}
117
118fn is_valid_bundle(dir: &Path) -> bool {
119    let weight = dir.join("weight.bin");
120    let config = dir.join("config.json");
121    if !weight.is_file() || !config.is_file() {
122        return false;
123    }
124    let Ok(raw) = std::fs::read_to_string(&config) else {
125        return false;
126    };
127    let Ok(v) = serde_json::from_str::<serde_json::Value>(&raw) else {
128        return false;
129    };
130    v.get("format")
131        .and_then(|x| x.as_str())
132        .map(|f| f == "aria-quant-bundle")
133        .unwrap_or(false)
134}
135
136/// Flatten a single top-level subdir (when config.json sits inside it).
137fn flatten_single_subdir(dir: &Path) -> std::io::Result<()> {
138    let mut entries: Vec<_> = std::fs::read_dir(dir)?
139        .filter_map(|e| e.ok())
140        .map(|e| e.path())
141        .collect();
142    if entries.len() != 1 {
143        return Ok(());
144    }
145    let only = &entries[0];
146    if !only.is_dir() {
147        return Ok(());
148    }
149    if !only.join("config.json").is_file() {
150        return Ok(());
151    }
152    let tmp = dir.join(format!(".flatten_{}", std::process::id()));
153    std::fs::rename(only, &tmp)?;
154    for entry in std::fs::read_dir(&tmp)? {
155        let entry = entry?;
156        std::fs::rename(entry.path(), dir.join(entry.file_name()))?;
157    }
158    std::fs::remove_dir_all(&tmp)?;
159    entries.clear();
160    Ok(())
161}
162
163fn extract_zip(data: &[u8], dest: &Path) -> Result<(), DownloadError> {
164    if data.len() < 4 || &data[0..2] != b"PK" {
165        return Err(DownloadError::BadZip("missing PK magic".into()));
166    }
167    std::fs::create_dir_all(dest)?;
168    let cursor = std::io::Cursor::new(data);
169    let mut archive = zip::ZipArchive::new(cursor)
170        .map_err(|e| DownloadError::BadZip(e.to_string()))?;
171    for i in 0..archive.len() {
172        let mut file = archive
173            .by_index(i)
174            .map_err(|e| DownloadError::Extract(e.to_string()))?;
175        let name = match file.enclosed_name() {
176            Some(n) => n.to_path_buf(),
177            None => continue,
178        };
179        let out_path = dest.join(&name);
180        if file.is_dir() {
181            std::fs::create_dir_all(&out_path)?;
182        } else {
183            if let Some(parent) = out_path.parent() {
184                std::fs::create_dir_all(parent)?;
185            }
186            let mut buf = Vec::with_capacity(file.size() as usize);
187            file.read_to_end(&mut buf)?;
188            let mut out = std::fs::File::create(&out_path)?;
189            out.write_all(&buf)?;
190        }
191    }
192    flatten_single_subdir(dest).map_err(DownloadError::Io)?;
193    Ok(())
194}
195
196fn atomic_replace(src: &Path, dst: &Path) -> std::io::Result<()> {
197    if dst.exists() {
198        std::fs::remove_dir_all(dst)?;
199    }
200    std::fs::rename(src, dst)
201}
202
203/// Download `model` from the Dashboard private source into
204/// `~/.ariacompute/models/{model}`, then return that directory.
205///
206/// If a valid bundle already exists at the cache path, the download is skipped.
207pub fn download_model(
208    model: &str,
209    token: &str,
210    site: Option<&str>,
211) -> Result<PathBuf, DownloadError> {
212    let (slug, quant) = parse_bundle_name(model)?;
213    let site = site.unwrap_or(DEFAULT_SITE);
214    let cache = models_dir()?.join(model);
215
216    if cache.exists() && is_valid_bundle(&cache) {
217        return Ok(cache);
218    }
219
220    let url = meta_url(site, &slug, &quant);
221    let agent = ureq::AgentBuilder::new().build();
222    let meta: DashboardMeta = agent
223        .get(&url)
224        .set("Authorization", &format!("Bearer {token}"))
225        .call()
226        .map_err(|e| DownloadError::Request(e.to_string()))?
227        .into_json::<DashboardMeta>()
228        .map_err(|e| DownloadError::Request(e.to_string()))?;
229    if meta.url.is_empty() {
230        return Err(DownloadError::Request(
231            "dashboard meta returned empty url".into(),
232        ));
233    }
234
235    let mut reader = agent
236        .get(&meta.url)
237        .set("Authorization", &format!("Bearer {token}"))
238        .call()
239        .map_err(|e| DownloadError::Stream(e.to_string()))?
240        .into_reader();
241    let mut data = Vec::new();
242    reader
243        .read_to_end(&mut data)
244        .map_err(|e| DownloadError::Stream(e.to_string()))?;
245
246    let staging = models_dir()?.join(format!(".{}.partial", model));
247    if staging.exists() {
248        std::fs::remove_dir_all(&staging)?;
249    }
250    extract_zip(&data, &staging)?;
251    if !is_valid_bundle(&staging) {
252        let _ = std::fs::remove_dir_all(&staging);
253        return Err(DownloadError::InvalidBundle(
254            "downloaded archive did not contain a valid aria-quant-bundle".into(),
255        ));
256    }
257    atomic_replace(&staging, &cache)?;
258    Ok(cache)
259}
aria_engine/download.rs

aria_engine/
download.rs
