super-stt-indexer/src/main.rs
Line | Count | Source |
1 | | // SPDX-License-Identifier: GPL-3.0-only |
2 | | //! `super-stt-indexer` — top-level orchestration. |
3 | | |
4 | | use std::path::PathBuf; |
5 | | |
6 | | use anyhow::Context; |
7 | | use clap::Parser; |
8 | | use log::{error, info, warn}; |
9 | | |
10 | | use super_stt_forge::{ForgeClient, RepoRef}; |
11 | | |
12 | | mod assets; |
13 | | mod carryforward; |
14 | | mod index_json; |
15 | | mod license; |
16 | | mod local; |
17 | | mod manifest; |
18 | | mod registry_toml; |
19 | | mod resolve; |
20 | | |
21 | | #[derive(Parser, Debug)] |
22 | | #[command(version, about)] |
23 | | struct Args { |
24 | | #[command(subcommand)] |
25 | | command: Command, |
26 | | } |
27 | | |
28 | | #[derive(clap::Subcommand, Debug)] |
29 | | enum Command { |
30 | | /// Build the published index from `registry.toml` + GitHub releases. |
31 | | Build(BuildArgs), |
32 | | /// Build a local index from staged backends — offline, no GitHub. For |
33 | | /// testing the daemon's download/install pipeline against a localhost |
34 | | /// static server. |
35 | | Local(local::LocalArgs), |
36 | | } |
37 | | |
38 | | #[derive(clap::Args, Debug)] |
39 | | struct BuildArgs { |
40 | | /// Path to `registry.toml` to read. |
41 | | #[arg(long, default_value = "registry/registry.toml")] |
42 | | registry: PathBuf, |
43 | | /// Path to the previously-published `index.json` (for carry-forward). If |
44 | | /// missing, falls through cleanly — bootstrap mode. |
45 | | #[arg(long)] |
46 | | prior_index: Option<PathBuf>, |
47 | | /// Where to write the new `index.json`. |
48 | | #[arg(long, default_value = "index.json")] |
49 | | out: PathBuf, |
50 | | } |
51 | | |
52 | | pub struct BuildFailure { |
53 | | pub error: String, |
54 | | pub attempted_version: Option<String>, |
55 | | pub attempted_tag: Option<String>, |
56 | | } |
57 | | |
58 | | #[tokio::main(flavor = "multi_thread", worker_threads = 4)] |
59 | 12 | async fn main() -> anyhow::Result<()> { |
60 | | // Workspace reqwest uses rustls without a bundled provider; install one. |
61 | 12 | super_stt_forge::install_crypto_provider(); |
62 | 12 | env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); |
63 | 12 | match Args::parse().command { |
64 | 12 | Command::Build(args2 ) => run_build2 (args).await, |
65 | 12 | Command::Local(args10 ) => local::run10 (&args10 ), |
66 | 12 | } |
67 | 12 | } |
68 | | |
69 | | /// Build the published index from `registry.toml` + GitHub releases. |
70 | 2 | async fn run_build(args: BuildArgs) -> anyhow::Result<()> { |
71 | 2 | let registry_text = std::fs::read_to_string(&args.registry) |
72 | 2 | .with_context(|| format!0 ("reading {}", args.registry.display()0 ))?0 ; |
73 | 2 | let registry = registry_toml::Registry::parse(®istry_text)?0 ; |
74 | | |
75 | 2 | let prior = match args.prior_index.as_ref() { |
76 | 0 | Some(p) if p.exists() => { |
77 | 0 | let text = std::fs::read_to_string(p)?; |
78 | 0 | Some(serde_json::from_str::<index_json::Index>(&text)?) |
79 | | } |
80 | 2 | _ => None, |
81 | | }; |
82 | | |
83 | | // Downloads release assets (subprocess bundles can be multi-GB) — use the |
84 | | // shared download client, not a timeout-less `Client::new()` that could hang |
85 | | // forever on a stalled connection. |
86 | 2 | let http = super_stt_forge::http::download_client(); |
87 | 2 | let now_iso = chrono_now_iso(); |
88 | | |
89 | 2 | let mut out_backends: Vec<index_json::IndexBackend> = Vec::new(); |
90 | | |
91 | 2 | for (id, entry) in ®istry.0 { |
92 | 2 | if entry.removed { |
93 | 0 | info!("skip `{id}` — removed"); |
94 | 0 | continue; |
95 | 2 | } |
96 | 2 | let client = super_stt_forge::client(entry.forge); |
97 | | // A malformed `repo` string must not abort the whole build — route it |
98 | | // through the same per-entry carry-forward path every other failure uses |
99 | | // (Tier 1 #28), instead of `?`-propagating out of the loop. |
100 | 2 | let built = match RepoRef::parse(&entry.repo) { |
101 | 2 | Ok(repo) => build_entry(client.as_ref(), &http, id, entry, &repo).await, |
102 | 0 | Err(e) => Err(BuildFailure { |
103 | 0 | error: format!("invalid repo `{}`: {e}", entry.repo), |
104 | 0 | attempted_version: None, |
105 | 0 | attempted_tag: None, |
106 | 0 | }), |
107 | | }; |
108 | 2 | match built { |
109 | 2 | Ok(b) => out_backends.push(b), |
110 | 0 | Err(failure) => { |
111 | 0 | error!("entry `{id}` failed: {}", failure.error); |
112 | 0 | let prior_entry = prior |
113 | 0 | .as_ref() |
114 | 0 | .and_then(|p| p.backends.iter().find(|b| b.id == *id)); |
115 | 0 | if let Some(carried) = carryforward::maybe_carry_forward( |
116 | 0 | id, |
117 | 0 | prior_entry, |
118 | 0 | &failure.error, |
119 | 0 | failure.attempted_version.as_deref().unwrap_or(""), |
120 | 0 | failure.attempted_tag.as_deref().unwrap_or(""), |
121 | 0 | &now_iso, |
122 | 0 | carryforward::MAX_STALENESS_DAYS, |
123 | 0 | ) { |
124 | 0 | warn!( |
125 | | "entry `{id}` — carrying forward last-known-good (v{})", |
126 | | carried.version |
127 | | ); |
128 | 0 | out_backends.push(carried); |
129 | 0 | } |
130 | | } |
131 | | } |
132 | | } |
133 | | |
134 | 2 | ensure_unique_sources(&out_backends)?0 ; |
135 | | |
136 | 2 | let index = index_json::Index { |
137 | 2 | schema_version: index_json::SCHEMA_VERSION, |
138 | 2 | generated_at: now_iso, |
139 | 2 | min_client: index_json::MIN_CLIENT.into(), |
140 | 2 | backends: out_backends, |
141 | 2 | }; |
142 | 2 | let text = serde_json::to_string_pretty(&index)?0 + "\n"; |
143 | 2 | super_stt_registry_types::fs::write_atomic(&args.out, text.as_bytes()) |
144 | 2 | .with_context(|| format!0 ("writing {}", args.out.display()0 ))?0 ; |
145 | 2 | info!( |
146 | | "wrote {} ({} backends)", |
147 | 2 | args.out.display(), |
148 | 2 | index.backends.len() |
149 | | ); |
150 | 2 | Ok(()) |
151 | 2 | } |
152 | | |
153 | 2 | async fn build_entry( |
154 | 2 | client: &dyn ForgeClient, |
155 | 2 | http: &reqwest::Client, |
156 | 2 | id: &str, |
157 | 2 | entry: ®istry_toml::Entry, |
158 | 2 | repo: &RepoRef, |
159 | 2 | ) -> Result<index_json::IndexBackend, BuildFailure> { |
160 | 2 | let resolved = resolve::resolve(client, repo, entry) |
161 | 2 | .await |
162 | 2 | .map_err(|e| BuildFailure { |
163 | 0 | error: format!("{e:#}"), |
164 | 0 | attempted_version: None, |
165 | 0 | attempted_tag: None, |
166 | 0 | })?; |
167 | | // From here the version + tag are known; record them on every later failure |
168 | | // so the carry-forward path can report what it tried to build. |
169 | 2 | let attempted_version = Some(resolved.version.to_string()); |
170 | 2 | let attempted_tag = Some(resolved.tag.clone()); |
171 | 2 | let fail = |e: &dyn std::fmt::Display| BuildFailure { |
172 | 0 | error: format!("{e:#}"), |
173 | 0 | attempted_version: attempted_version.clone(), |
174 | 0 | attempted_tag: attempted_tag.clone(), |
175 | 0 | }; |
176 | | |
177 | | // The manifest is the `backend.toml` release asset: parse + validate the |
178 | | // exact bytes that get hashed, so reviewed == pinned == installed. A release |
179 | | // without the asset is not installable (no synthesize fallback) — fail the |
180 | | // entry. |
181 | 2 | let (url, _declared_size) = |
182 | 2 | assets::resolve_url("backend.toml", &resolved.release.assets).map_err(|e| fail0 (&e0 ))?0 ; |
183 | 2 | let (bytes, sha256) = assets::fetch_manifest_asset(http, &url) |
184 | 2 | .await |
185 | 2 | .map_err(|e| fail0 (&e0 ))?0 ; |
186 | 2 | let size = bytes.len() as u64; |
187 | 2 | let text = String::from_utf8(bytes).map_err(|e| fail0 (&e0 ))?0 ; |
188 | 2 | let m = manifest::Manifest::parse(&text).map_err(|e| fail0 (&e0 ))?0 ; |
189 | 2 | let manifest_pin = Some(index_json::IndexAsset { url, size, sha256 }); |
190 | 2 | manifest::validate(&m, &resolved.version, &entry.repo).map_err(|e| fail0 (&e0 ))?0 ; |
191 | 2 | let idx_assets = resolve_index_assets(http, &m, &resolved.release.assets) |
192 | 2 | .await |
193 | 2 | .map_err(|e| fail0 (&e0 ))?0 ; |
194 | | |
195 | 2 | Ok(into_index_backend( |
196 | 2 | id, |
197 | 2 | m, |
198 | 2 | resolved.version.to_string(), |
199 | 2 | resolved.tag, |
200 | 2 | idx_assets, |
201 | 2 | manifest_pin, |
202 | 2 | )) |
203 | 2 | } |
204 | | |
205 | | /// A unique temp path for staging one downloaded asset part during validation. |
206 | 0 | fn temp_part_path() -> std::path::PathBuf { |
207 | | use std::sync::atomic::{AtomicU64, Ordering}; |
208 | | static N: AtomicU64 = AtomicU64::new(0); |
209 | 0 | let n = N.fetch_add(1, Ordering::Relaxed); |
210 | 0 | std::env::temp_dir().join(format!("stt-idx-{}-{n}.part", std::process::id())) |
211 | 0 | } |
212 | | |
213 | | /// RAII owner of the downloaded part files: removes them all on drop, so a |
214 | | /// mid-loop download/validation error can't leak the (possibly multi-GB) parts |
215 | | /// already fetched (Tier 1 #29). A path is registered *before* its download so |
216 | | /// even a partially-written part is cleaned up. |
217 | | struct TempParts(Vec<std::path::PathBuf>); |
218 | | |
219 | | impl TempParts { |
220 | 0 | fn new() -> Self { |
221 | 0 | Self(Vec::new()) |
222 | 0 | } |
223 | 0 | fn register(&mut self, path: std::path::PathBuf) { |
224 | 0 | self.0.push(path); |
225 | 0 | } |
226 | 0 | fn paths(&self) -> &[std::path::PathBuf] { |
227 | 0 | &self.0 |
228 | 0 | } |
229 | | } |
230 | | |
231 | | impl Drop for TempParts { |
232 | 0 | fn drop(&mut self) { |
233 | 0 | for p in &self.0 { |
234 | 0 | let _ = std::fs::remove_file(p); |
235 | 0 | } |
236 | 0 | } |
237 | | } |
238 | | |
239 | | /// Build the index entry for one subprocess variant from its downloaded part |
240 | | /// pins: a single-file pin (`url`/`size`/`sha256`) or a multi-part pin. |
241 | 0 | fn subprocess_index_entry( |
242 | 0 | sa: &manifest::SubprocessAsset, |
243 | 0 | mut pins: Vec<index_json::IndexAsset>, |
244 | 0 | ) -> index_json::IndexSubprocessAsset { |
245 | 0 | let (url, size, sha256, parts) = if sa.is_multipart() { |
246 | 0 | (None, None, None, pins) |
247 | | } else { |
248 | 0 | let p = pins.remove(0); |
249 | 0 | (Some(p.url), Some(p.size), Some(p.sha256), Vec::new()) |
250 | | }; |
251 | 0 | index_json::IndexSubprocessAsset { |
252 | 0 | target: sa.target.clone(), |
253 | 0 | accel: sa.accel.to_string(), |
254 | 0 | cuda_major: sa.cuda_major, |
255 | 0 | cuda_sm: sa.cuda_sm, |
256 | 0 | cudnn: sa.cudnn, |
257 | 0 | url, |
258 | 0 | size, |
259 | 0 | sha256, |
260 | 0 | parts, |
261 | 0 | } |
262 | 0 | } |
263 | | |
264 | | /// Resolve and hash the binary artifacts a release declares — the wasm |
265 | | /// component or each subprocess variant — into the index's asset block. |
266 | 2 | async fn resolve_index_assets( |
267 | 2 | http: &reqwest::Client, |
268 | 2 | m: &manifest::Manifest, |
269 | 2 | release_assets: &[super_stt_forge::ReleaseAsset], |
270 | 2 | ) -> anyhow::Result<index_json::IndexAssets> { |
271 | 2 | let mut idx_assets = index_json::IndexAssets::default(); |
272 | 2 | if let Some(wasm) = &m.assets.wasm { |
273 | 2 | let (url, size) = assets::resolve_url(wasm, release_assets)?0 ; |
274 | 2 | let sha = assets::fetch_wasm_and_hash(http, &url, wasm).await?0 ; |
275 | 2 | idx_assets.wasm = Some(index_json::IndexAsset { |
276 | 2 | url, |
277 | 2 | size, |
278 | 2 | sha256: sha, |
279 | 2 | }); |
280 | 0 | } |
281 | 2 | for sa0 in &m.assets.subprocess { |
282 | | // Download each part (a single-file variant has one) to a temp file, |
283 | | // hashing it, then validate the reassembled archive before pinning. |
284 | | // `TempParts` removes every downloaded part on scope exit — including on |
285 | | // an early `?` from `resolve_url`/`download_to_file`/validation — so a |
286 | | // mid-loop error can't leak multi-GB temp files (Tier 1 #29). |
287 | 0 | let files = sa.release_files(); |
288 | 0 | let mut tmp = TempParts::new(); |
289 | 0 | let mut pins: Vec<index_json::IndexAsset> = Vec::with_capacity(files.len()); |
290 | 0 | for f in &files { |
291 | 0 | let (url, _) = assets::resolve_url(f, release_assets)?; |
292 | 0 | let dest = temp_part_path(); |
293 | 0 | tmp.register(dest.clone()); |
294 | 0 | let (size, sha256) = assets::download_to_file(http, &url, f, &dest).await?; |
295 | 0 | pins.push(index_json::IndexAsset { url, size, sha256 }); |
296 | | } |
297 | 0 | assets::validate_subprocess_parts(tmp.paths(), &sa.label(), &m.backend.entrypoint)?; |
298 | 0 | idx_assets.subprocess.push(subprocess_index_entry(sa, pins)); |
299 | | // `tmp` drops here (validation done), removing the parts. |
300 | | } |
301 | 2 | Ok(idx_assets) |
302 | 2 | } |
303 | | |
304 | | /// Assemble the published `IndexBackend` from a validated manifest, its |
305 | | /// resolved `version` + `tag`, and the hashed assets. Thin wrapper over the |
306 | | /// canonical [`index_json::IndexBackend::from_manifest`] synthesis (shared with |
307 | | /// the daemon's Custom-repo and local-dir install paths) — the indexer supplies |
308 | | /// the maintainer-declared `id` rather than deriving it from `source`. |
309 | 14 | pub(crate) fn into_index_backend( |
310 | 14 | id: &str, |
311 | 14 | m: manifest::Manifest, |
312 | 14 | version: String, |
313 | 14 | tag: String, |
314 | 14 | assets: index_json::IndexAssets, |
315 | 14 | manifest: Option<index_json::IndexAsset>, |
316 | 14 | ) -> index_json::IndexBackend { |
317 | 14 | index_json::IndexBackend::from_manifest(id.to_string(), m, version, tag, assets, manifest) |
318 | 14 | } |
319 | | |
320 | | /// A backend's `source` is its unique identity. Two distinct entries that |
321 | | /// collide on `source` would be indistinguishable to every daemon (the |
322 | | /// daemon's `dedup_sources` guard silently drops all but the first), so a |
323 | | /// collision must never be published — fail the build instead. |
324 | 6 | fn ensure_unique_sources(backends: &[index_json::IndexBackend]) -> anyhow::Result<()> { |
325 | 6 | let mut seen: std::collections::HashMap<&str, &str> = std::collections::HashMap::new(); |
326 | 10 | for b in backends6 { |
327 | 10 | if let Some(prev_id2 ) = seen.insert(b.source.as_str(), b.id.as_str()) { |
328 | 2 | anyhow::bail!( |
329 | | "duplicate source `{}` shared by entries `{}` and `{}`; each backend must have a distinct source", |
330 | | b.source, |
331 | | prev_id, |
332 | | b.id, |
333 | | ); |
334 | 8 | } |
335 | | } |
336 | 4 | Ok(()) |
337 | 6 | } |
338 | | |
339 | 12 | fn chrono_now_iso() -> String { |
340 | 12 | chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string() |
341 | 12 | } |
342 | | |
343 | | #[cfg(test)] |
344 | | mod tests { |
345 | | use super::*; |
346 | | |
347 | 8 | fn backend(id: &str, source: &str) -> index_json::IndexBackend { |
348 | 8 | index_json::IndexBackend { |
349 | 8 | id: id.into(), |
350 | 8 | source: source.into(), |
351 | 8 | version: "1.0.0".into(), |
352 | 8 | tag: "v1.0.0".into(), |
353 | 8 | name: id.into(), |
354 | 8 | description: None, |
355 | 8 | license: "Apache-2.0".into(), |
356 | 8 | kind: "wasm".into(), |
357 | 8 | contract: "v1".into(), |
358 | 8 | entrypoint: format!("{id}.wasm"), |
359 | 8 | allowed_hosts: Vec::new(), |
360 | 8 | online: false, |
361 | 8 | supports_gpu: false, |
362 | 8 | supports_cpu: true, |
363 | 8 | models: Vec::new(), |
364 | 8 | secrets: Vec::new(), |
365 | 8 | options: Vec::new(), |
366 | 8 | assets: index_json::IndexAssets::default(), |
367 | 8 | index_stale: None, |
368 | 8 | manifest: None, |
369 | 8 | } |
370 | 8 | } |
371 | | |
372 | | #[test] |
373 | 2 | fn unique_sources_pass() { |
374 | 2 | let backends = vec![ |
375 | 2 | backend("mistral", "github.com/x/y/mistral"), |
376 | 2 | backend("openai", "github.com/x/y/openai"), |
377 | | ]; |
378 | 2 | ensure_unique_sources(&backends).unwrap(); |
379 | 2 | } |
380 | | |
381 | | #[test] |
382 | 2 | fn duplicate_sources_are_rejected() { |
383 | 2 | let backends = vec![ |
384 | 2 | backend("mistral", "github.com/x/y"), |
385 | 2 | backend("openai", "github.com/x/y"), |
386 | | ]; |
387 | 2 | let err = ensure_unique_sources(&backends).unwrap_err(); |
388 | 2 | assert!(err.to_string().contains("duplicate source")); |
389 | 2 | } |
390 | | } |