rlx_cli/
auto_dispatch.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Auto-dispatch: pick a registered model runner from a weights path.
17//!
18//! `auto_runner_name(path)` resolves the path (file or directory), sniffs
19//! the model family (GGUF `general.architecture` for `.gguf`, sidecar
20//! `config.json` `model_type` for safetensors), and maps it to the short
21//! runner name a callsite registered with [`register_cli`](crate::register_cli)
22//! (e.g. `"qwen3"`, `"gemma"`).
23//!
24//! `auto_dispatch(path, args)` is a one-shot: sniff, look up, run.
25//!
26//! Used by `skill` so callers don't need to hardcode `Qwen3Runner` vs
27//! `GemmaRunner` per family.
28
29use anyhow::{Context, Result, anyhow, bail};
30use rlx_core::gguf_config::{
31    DINOV2_GGUF_ARCHES, FLUX_GGUF_ARCHES, SAM_GGUF_ARCHES, SAM2_GGUF_ARCHES, SAM3_GGUF_ARCHES,
32    VJEPA2_GGUF_ARCHES, W2V_BERT_GGUF_ARCHES,
33};
34use rlx_core::gguf_support::{
35    gguf_architecture_from_path, gguf_family_for_arch, resolve_weights_file,
36};
37use std::path::{Path, PathBuf};
38
39use crate::registry::run_registered;
40
41/// Entry point for an `rlx-run auto WEIGHTS [args...]` subcommand.
42///
43/// Treats the first positional as the weights path (file or directory),
44/// sniffs the runner, and forwards the remaining args to it. The
45/// canonical wiring is `register_cli("auto", "...", rlx_cli::run_auto)`
46/// in the multiplexer.
47pub fn run_auto(args: &[String]) -> Result<()> {
48    let Some(first) = args.first() else {
49        bail!(
50            "auto: expected WEIGHTS path as the first argument\n\
51             usage: rlx-run auto <weights-path> [runner-args...]"
52        );
53    };
54    if matches!(first.as_str(), "-h" | "--help" | "help") {
55        println!(
56            "rlx-run auto — sniff a GGUF / safetensors file and dispatch to the right runner\n\
57             \n\
58             USAGE:\n  rlx-run auto <weights-path> [runner-args...]\n\
59             \n\
60             The first argument is forwarded as the runner's --weights value;\n\
61             remaining arguments are passed through unchanged."
62        );
63        return Ok(());
64    }
65    let path = Path::new(first);
66    let sniff = auto_sniff(path)?;
67    eprintln!(
68        "[rlx-run auto] {} → runner `{}` (from {:?})",
69        sniff.path.display(),
70        sniff.runner_name,
71        sniff.from
72    );
73    // Re-build argv: most per-family runners take `--weights PATH`. If the
74    // caller already passed --weights, don't double it; otherwise inject.
75    let rest: Vec<String> = args[1..].to_vec();
76    let has_weights_flag = rest
77        .iter()
78        .any(|a| a == "--weights" || a.starts_with("--weights="));
79    let mut forwarded: Vec<String> = Vec::with_capacity(rest.len() + 2);
80    if !has_weights_flag {
81        forwarded.push("--weights".into());
82        forwarded.push(sniff.path.display().to_string());
83    }
84    forwarded.extend(rest);
85    match run_registered(sniff.runner_name, &forwarded)? {
86        Some(()) => Ok(()),
87        None => bail!(
88            "auto: runner `{}` not registered (sniffed from {:?}); register it via \
89             `register_cli` in your binary's main",
90            sniff.runner_name,
91            sniff.from
92        ),
93    }
94}
95
96/// Source the sniffer used to identify the model family.
97#[derive(Debug, Clone)]
98pub enum SniffedFrom {
99    /// `general.architecture` value read from a `.gguf` file.
100    GgufArch(String),
101    /// `model_type` value read from a sidecar `config.json`.
102    SafetensorsConfig(String),
103}
104
105/// Result of sniffing a weights path.
106#[derive(Debug, Clone)]
107pub struct SniffedRunner {
108    /// Concrete file we sniffed (after resolving a directory).
109    pub path: PathBuf,
110    /// Short runner name as registered with `register_cli`.
111    pub runner_name: &'static str,
112    /// Where the sniff came from — useful for diagnostics.
113    pub from: SniffedFrom,
114}
115
116/// A catalog arch that RLX recognizes but has not yet implemented a runner
117/// for. Returned by [`known_unimplemented_arch`] so error messages can point
118/// at the PLAN.md milestone that unblocks the family.
119#[derive(Debug, Clone, Copy, PartialEq, Eq)]
120pub struct UnimplementedArch {
121    /// Display name (e.g. `"Mistral 3.5"`).
122    pub family: &'static str,
123    /// PLAN.md milestone tag (e.g. `"M4"`).
124    pub milestone: &'static str,
125    /// One-line note for the user.
126    pub note: &'static str,
127}
128
129/// Family-level metadata referenced by [`KNOWN_UNIMPLEMENTED`]. Static so
130/// the phf map can hold `&'static UnimplementedArch`.
131mod families {
132    use super::UnimplementedArch;
133    pub static MISTRAL: UnimplementedArch = UnimplementedArch {
134        family: "Mistral 3+ / Ministral",
135        milestone: "M4",
136        note: "Llama-shaped with newer RoPE; share `rlx-llama-base` per PLAN.md M4",
137    };
138    pub static PHI: UnimplementedArch = UnimplementedArch {
139        family: "Phi 3 / Phi 4",
140        milestone: "M4",
141        note: "Phi3/4 share llama.cpp arch tag — PLAN.md M4",
142    };
143    pub static PHIMOE: UnimplementedArch = UnimplementedArch {
144        family: "Phi MoE",
145        milestone: "M4 + M5",
146        note: "Phi + MoE routing; depends on shared MoE block — PLAN.md M4/M5",
147    };
148    pub static BONSAI: UnimplementedArch = UnimplementedArch {
149        family: "Bonsai",
150        milestone: "M4",
151        note: "Llama-shaped; HF model_type only — usually ships as llama GGUF — PLAN.md M4",
152    };
153    pub static OMNICODER: UnimplementedArch = UnimplementedArch {
154        family: "OmniCoder",
155        milestone: "M4",
156        note: "Qwen3-coder shaped — PLAN.md M4 (often tagged `qwen3` in GGUF)",
157    };
158    pub static MINIMAX: UnimplementedArch = UnimplementedArch {
159        family: "MiniMax M2",
160        milestone: "M5",
161        note: "Lightning Attention; depends on `rlx-ssm` upstream — PLAN.md M5",
162    };
163    pub static GLM: UnimplementedArch = UnimplementedArch {
164        family: "GLM 4 / 5",
165        milestone: "M5",
166        note: "GLM RoPE + RMSNorm placement — PLAN.md M5",
167    };
168    pub static GLM_MOE: UnimplementedArch = UnimplementedArch {
169        family: "GLM 4 MoE",
170        milestone: "M5",
171        note: "GLM + MoE routing — PLAN.md M5",
172    };
173    pub static GPT_OSS: UnimplementedArch = UnimplementedArch {
174        family: "gpt-oss",
175        milestone: "M5",
176        note: "OpenAI gpt-oss — confirm arch shape — PLAN.md M5",
177    };
178    pub static NEMOTRON: UnimplementedArch = UnimplementedArch {
179        family: "Nemotron",
180        milestone: "M5",
181        note: "Dense Nemotron arch — PLAN.md M5",
182    };
183    pub static NEMOTRON_H: UnimplementedArch = UnimplementedArch {
184        family: "Nemotron-H",
185        milestone: "M5",
186        note: "Mamba+attention hybrid; depends on `rlx-ssm` upstream — PLAN.md M5/M7",
187    };
188    #[allow(dead_code)]
189    pub static LFM: UnimplementedArch = UnimplementedArch {
190        family: "LFM 2 / 2.5",
191        milestone: "M5",
192        note: "Liquid Foundation Models with custom SSM layers — PLAN.md M5",
193    };
194    pub static LFM_MOE: UnimplementedArch = UnimplementedArch {
195        family: "LFM 2 MoE",
196        milestone: "M5",
197        note: "LFM + MoE — PLAN.md M5",
198    };
199    pub static QWEN3_MOE: UnimplementedArch = UnimplementedArch {
200        family: "Qwen3 MoE",
201        milestone: "M5",
202        note: "Qwen3 + MoE routing block — PLAN.md M5 (often loadable via qwen3 runner once MoE lands)",
203    };
204    pub static QWEN3_NEXT: UnimplementedArch = UnimplementedArch {
205        family: "Qwen3-Next",
206        milestone: "M5",
207        note: "Qwen3-Next variant — confirm arch deltas vs qwen3 — PLAN.md M5",
208    };
209    pub static GEMMA3: UnimplementedArch = UnimplementedArch {
210        family: "Gemma 3",
211        milestone: "M2",
212        note: "Gemma 3 (270m / 4b / 12b / 27b) adds per-layer sliding window + new RoPE — \
213               needs rlx-gemma config branch — PLAN.md M2",
214    };
215    pub static GEMMA3N: UnimplementedArch = UnimplementedArch {
216        family: "Gemma 3n",
217        milestone: "M2",
218        note: "Gemma 3n (mobile/edge Matformer variant) — PLAN.md M2",
219    };
220    pub static GEMMA4: UnimplementedArch = UnimplementedArch {
221        family: "Gemma 4 MoE (A4B)",
222        milestone: "M2",
223        note: "Gemma 4 MoE A4B routing block — dense + E2B/E4B run via the `gemma` runner — PLAN.md M2",
224    };
225    pub static QWEN3_VL: UnimplementedArch = UnimplementedArch {
226        family: "Qwen3-VL",
227        milestone: "M7",
228        note: "vision tower + projector + LM (dense or MoE) — PLAN.md M7",
229    };
230    pub static QWEN3_MTP: UnimplementedArch = UnimplementedArch {
231        family: "Qwen3 / Qwen3.6 + MTP",
232        milestone: "M6",
233        note: "multi-token-prediction draft heads — PLAN.md M6",
234    };
235    pub static LLADA: UnimplementedArch = UnimplementedArch {
236        family: "LLaDA / LLaDA MoE (text-only)",
237        milestone: "M5",
238        note: "dense LLaDA arch in llama.cpp; rlx-llada2 currently targets the diffusion runner — PLAN.md M5",
239    };
240    pub static GRANITE: UnimplementedArch = UnimplementedArch {
241        family: "Granite (IBM)",
242        milestone: "M4",
243        note: "Llama-shaped — PLAN.md M4",
244    };
245    pub static DEEPSEEK: UnimplementedArch = UnimplementedArch {
246        family: "DeepSeek 2",
247        milestone: "M5",
248        note: "MoE + MLA attention — needs MoE block + MLA primitive — PLAN.md M5",
249    };
250    pub static COHERE: UnimplementedArch = UnimplementedArch {
251        family: "Command-R / Cohere",
252        milestone: "M4",
253        note: "Llama-shaped — PLAN.md M4",
254    };
255}
256
257/// Catalog families we know about but haven't implemented yet.
258///
259/// The keys are the **actual** GGUF `general.architecture` strings llama.cpp
260/// uses (`src/llama-arch.cpp::LLM_ARCH_NAMES`) plus their HF `model_type`
261/// aliases when those differ. Notably:
262///
263/// * Mistral 1/2 and Qwen 2.5 ship as `general.architecture = llama` /
264///   `qwen2` respectively — they don't have their own llama.cpp arch tag.
265///   Those tags route to the existing `llama32` / `qwen3` runners and are
266///   *not* listed here.
267/// * Mistral 3+ ships as `mistral3` / `mistral4` (real tags).
268/// * Phi-4 ships as `phi3` (Phi-4 reuses the Phi-3 arch in llama.cpp).
269///
270/// Both GGUF arch tags and HF `model_type` values are accepted so
271/// downstream callers don't keep two parallel lists.
272static KNOWN_UNIMPLEMENTED: phf::Map<&'static str, &'static UnimplementedArch> = phf::phf_map! {
273    // Mistral / Ministral (real llama.cpp tags)
274    "mistral3" => &families::MISTRAL,
275    "mistral4" => &families::MISTRAL,
276    // Phi family — Llama32Family accepts the arch tag, but the GGUF
277    // tensor-name remap for `phi3`/`phi4` (e.g. `blk.*.attn_q.weight`
278    // → `model.layers.*.self_attn.q_proj.weight`) is M4 follow-up.
279    "phi3" => &families::PHI,
280    "phi4" => &families::PHI,
281    "phimoe" => &families::PHIMOE,
282    // Catalog HF model_type aliases — same remap gap as phi3.
283    "bonsai" => &families::BONSAI,
284    "omnicoder" => &families::OMNICODER,
285    // Hybrid / SSM families
286    "minimax-m2" => &families::MINIMAX,
287    "minimax_m2" => &families::MINIMAX,
288    "minimax" => &families::MINIMAX,
289    "glm4" => &families::GLM,
290    "glm5" => &families::GLM,
291    "chatglm" => &families::GLM,
292    "glm4moe" => &families::GLM_MOE,
293    "gpt-oss" => &families::GPT_OSS,
294    "gpt_oss" => &families::GPT_OSS,
295    "nemotron" => &families::NEMOTRON,
296    "nemotron_h" => &families::NEMOTRON_H,
297    "nemotron_h_moe" => &families::NEMOTRON_H,
298    // lfm2 / lfm / lfm25 / lfm2_5 are now routed through `rlx-lfm`'s
299    // `LfmRunner` via `gguf_family_for_arch` → `GgufModelFamily::Lfm`.
300    // Only the MoE variant remains unimplemented.
301    "lfm2moe" => &families::LFM_MOE,
302    // Qwen variants we don't run yet
303    "qwen3moe" => &families::QWEN3_MOE,
304    "qwen3next" => &families::QWEN3_NEXT,
305    // Gemma 3 / 3n still pending; Gemma 4 dense + E2B now route to the
306    // `gemma` runner (see `model_type_runner_name` / `arch_runner_name`).
307    // Only the MoE A4B variant remains unimplemented.
308    "gemma3" => &families::GEMMA3,
309    "gemma3n" => &families::GEMMA3N,
310    "gemma4moe" => &families::GEMMA4,
311    "qwen3vl" => &families::QWEN3_VL,
312    "qwen3vlmoe" => &families::QWEN3_VL,
313    "qwen3_vl" => &families::QWEN3_VL,
314    "qwen3-vl" => &families::QWEN3_VL,
315    "qwen3_mtp" => &families::QWEN3_MTP,
316    "qwen3-mtp" => &families::QWEN3_MTP,
317    "qwen36_mtp" => &families::QWEN3_MTP,
318    // Other catalog-adjacent families
319    "llada" => &families::LLADA,
320    "llada-moe" => &families::LLADA,
321    "granite" => &families::GRANITE,
322    "granitemoe" => &families::GRANITE,
323    "granitehybrid" => &families::GRANITE,
324    "deepseek2" => &families::DEEPSEEK,
325    "deepseek2-ocr" => &families::DEEPSEEK,
326    "command-r" => &families::COHERE,
327    "cohere2" => &families::COHERE,
328};
329
330/// Look up an arch / model_type in the unimplemented-families table.
331pub fn known_unimplemented_arch(arch_or_model_type: &str) -> Option<UnimplementedArch> {
332    KNOWN_UNIMPLEMENTED.get(arch_or_model_type).map(|p| **p)
333}
334
335/// Snapshot of every (key, family) pair currently in the unimplemented
336/// table — useful for `rlx-run check --list-unimplemented` style tooling.
337pub fn known_unimplemented_keys() -> impl Iterator<Item = (&'static str, &'static UnimplementedArch)>
338{
339    KNOWN_UNIMPLEMENTED.entries().map(|(k, v)| (*k, *v))
340}
341
342/// Map a GGUF `general.architecture` tag to the short runner name.
343///
344/// Returns `None` for embed-only families (`bert`, `nomic-bert`, …) which
345/// aren't currently exposed through the `rlx-run` dispatch table, and for
346/// catalog families that aren't implemented yet — those get a richer error
347/// via [`known_unimplemented_arch`] when sniffed.
348pub fn arch_runner_name(arch: &str) -> Option<&'static str> {
349    if let Some(fam) = gguf_family_for_arch(arch) {
350        return Some(fam.runner_name());
351    }
352    if FLUX_GGUF_ARCHES.contains(&arch) {
353        return Some("flux2");
354    }
355    if DINOV2_GGUF_ARCHES.contains(&arch) {
356        return Some("dinov2");
357    }
358    if VJEPA2_GGUF_ARCHES.contains(&arch) {
359        return Some("vjepa2");
360    }
361    if SAM3_GGUF_ARCHES.contains(&arch) {
362        return Some("sam3");
363    }
364    if SAM2_GGUF_ARCHES.contains(&arch) {
365        return Some("sam2");
366    }
367    if SAM_GGUF_ARCHES.contains(&arch) {
368        return Some("sam1");
369    }
370    if W2V_BERT_GGUF_ARCHES.contains(&arch) {
371        return Some("wav2vec2-bert");
372    }
373    None
374}
375
376/// Map an HF `config.json` `model_type` value to a short runner name.
377///
378/// HF naming differs from GGUF tags — `model_type: "llama"` covers Llama
379/// 2 / 3 / 3.x, `qwen3` covers Qwen3 and Qwen3 MoE, etc.
380pub fn model_type_runner_name(model_type: &str) -> Option<&'static str> {
381    match model_type {
382        // qwen2 deliberately omitted — rlx-qwen3 doesn't support
383        // Qwen 2 tensor layout (needs q/k/v bias + no QK-norm).
384        // qwen2 GGUFs fall through to known_unimplemented_arch.
385        "qwen3" | "qwen3_moe" | "qwen3moe" | "qwen25" | "qwen2_5" | "qwen2.5" | "qwen251"
386        | "qwen2_5_1" => Some("qwen3"),
387        "qwen35" | "qwen3_5" | "qwen35_moe" | "qwen35moe" => Some("qwen35"),
388        // Qwen3.6 runs through the qwen35 trunk (PLAN.md M1).
389        "qwen36" | "qwen3_6" | "qwen36_moe" | "qwen36moe" => Some("qwen35"),
390        "llama" | "llama2" | "llama3" => Some("llama32"),
391        // Gemma 4 dense + E2B/E4B mobile (PLE + KV-shared). The unified
392        // top-level model_type is `gemma4`; the text sub-config is
393        // `gemma4_text`. The MoE A4B variant (`gemma4moe`) stays in the
394        // unimplemented table until its routing block is validated.
395        "gemma"
396        | "gemma2"
397        | "gemma3"
398        | "gemma3n"
399        | "gemma4"
400        | "gemma4_text"
401        | "gemma4_unified"
402        | "gemma4_unified_text" => Some("gemma"),
403        "dinov2" | "dinov2_with_registers" => Some("dinov2"),
404        "vjepa2" | "vjepa" => Some("vjepa2"),
405        "sam" | "sam_vit" | "mobile-sam" | "mobile_sam" => Some("sam1"),
406        "sam2" => Some("sam2"),
407        "sam3" => Some("sam3"),
408        "whisper" => Some("whisper"),
409        "wav2vec2-bert" | "wav2vec2_bert" | "w2v-bert" | "w2v_bert" => Some("wav2vec2-bert"),
410        "flux" | "flux2" => Some("flux2"),
411        _ => None,
412    }
413}
414
415/// Sniff `model_type` from the `config.json` next to a safetensors file.
416fn read_model_type_from_sidecar(path: &Path) -> Result<Option<String>> {
417    let dir = path
418        .parent()
419        .ok_or_else(|| anyhow!("safetensors path {path:?} has no parent dir"))?;
420    let cfg = dir.join("config.json");
421    if !cfg.is_file() {
422        return Ok(None);
423    }
424    let bytes = std::fs::read(&cfg).with_context(|| format!("reading {cfg:?}"))?;
425    let v: serde_json::Value =
426        serde_json::from_slice(&bytes).with_context(|| format!("parsing {cfg:?}"))?;
427    Ok(v.get("model_type")
428        .and_then(serde_json::Value::as_str)
429        .map(str::to_owned))
430}
431
432/// Resolve `path` to a single weight file, then sniff the runner.
433pub fn auto_sniff(path: &Path) -> Result<SniffedRunner> {
434    let file = resolve_weights_file(path)?;
435    let ext = file.extension().and_then(|s| s.to_str()).unwrap_or("");
436    match ext {
437        "gguf" => {
438            let arch = gguf_architecture_from_path(&file)?;
439            let runner = arch_runner_name(&arch).ok_or_else(|| {
440                if let Some(u) = known_unimplemented_arch(&arch) {
441                    anyhow!(
442                        "{file:?}: GGUF architecture `{arch}` is {} ({}) — not yet implemented in rlx-models. {}",
443                        u.family, u.milestone, u.note
444                    )
445                } else {
446                    anyhow!(
447                        "{file:?}: GGUF architecture `{arch}` has no registered rlx runner; \
448                         see `rlx-run` for supported families"
449                    )
450                }
451            })?;
452            Ok(SniffedRunner {
453                path: file,
454                runner_name: runner,
455                from: SniffedFrom::GgufArch(arch),
456            })
457        }
458        "safetensors" => {
459            let model_type = read_model_type_from_sidecar(&file)?.ok_or_else(|| {
460                anyhow!("{file:?}: no `model_type` in sidecar config.json (auto-dispatch needs it)")
461            })?;
462            let runner = model_type_runner_name(&model_type).ok_or_else(|| {
463                if let Some(u) = known_unimplemented_arch(&model_type) {
464                    anyhow!(
465                        "{file:?}: safetensors model_type `{model_type}` is {} ({}) — not yet implemented in rlx-models. {}",
466                        u.family, u.milestone, u.note
467                    )
468                } else {
469                    anyhow!(
470                        "{file:?}: safetensors model_type `{model_type}` has no registered rlx runner"
471                    )
472                }
473            })?;
474            Ok(SniffedRunner {
475                path: file,
476                runner_name: runner,
477                from: SniffedFrom::SafetensorsConfig(model_type),
478            })
479        }
480        other => {
481            bail!("{file:?}: unsupported extension `.{other}` (expected .gguf or .safetensors)")
482        }
483    }
484}
485
486/// Sniff `path` and return only the runner short name.
487pub fn auto_runner_name(path: &Path) -> Result<&'static str> {
488    Ok(auto_sniff(path)?.runner_name)
489}
490
491/// Sniff `path`, look up its runner in the registry, and run it with `args`.
492///
493/// `args` should be the per-runner argv *without* the leading subcommand.
494/// Returns the runner name that was dispatched to.
495pub fn auto_dispatch(path: &Path, args: &[String]) -> Result<&'static str> {
496    let sniff = auto_sniff(path)?;
497    match run_registered(sniff.runner_name, args)? {
498        Some(()) => Ok(sniff.runner_name),
499        None => bail!(
500            "runner `{}` not registered (sniffed from {:?}); register it via \
501             `register_cli` before calling auto_dispatch",
502            sniff.runner_name,
503            sniff.from
504        ),
505    }
506}
507
508#[cfg(test)]
509mod tests {
510    use super::*;
511
512    #[test]
513    fn arch_runner_maps_lm_families() {
514        assert_eq!(arch_runner_name("qwen3"), Some("qwen3"));
515        // qwen2 now routes to the qwen3 runner — the runner reads
516        // attention_bias + qk_norm from the GGUF arch tag and emits
517        // the right per-layer math.
518        assert_eq!(arch_runner_name("qwen2"), Some("qwen3"));
519        assert_eq!(arch_runner_name("qwen35"), Some("qwen35"));
520        assert_eq!(arch_runner_name("qwen35moe"), Some("qwen35"));
521        // Qwen3.6 reuses the qwen35 trunk (PLAN.md M1). qwen36_mtp still
522        // routes through known_unimplemented_arch — base qwen36 routes
523        // here so unsloth/Qwen3.6-27B-GGUF (no MTP) just works.
524        assert_eq!(arch_runner_name("qwen36"), Some("qwen35"));
525        assert_eq!(arch_runner_name("qwen36moe"), Some("qwen35"));
526        // Qwen 2.5 / 2.5.1 ship as `qwen2` arch tag; explicit short
527        // tags also route to the qwen3 runner (PLAN.md M4).
528        assert_eq!(arch_runner_name("qwen25"), Some("qwen3"));
529        assert_eq!(arch_runner_name("qwen2_5"), Some("qwen3"));
530        assert_eq!(arch_runner_name("llama"), Some("llama32"));
531        assert_eq!(arch_runner_name("gemma"), Some("gemma"));
532        assert_eq!(arch_runner_name("gemma2"), Some("gemma"));
533    }
534
535    #[test]
536    fn arch_runner_maps_vision_and_diffusion() {
537        assert_eq!(arch_runner_name("dinov2"), Some("dinov2"));
538        assert_eq!(arch_runner_name("sam"), Some("sam1"));
539        assert_eq!(arch_runner_name("mobile-sam"), Some("sam1"));
540        assert_eq!(arch_runner_name("sam2"), Some("sam2"));
541        assert_eq!(arch_runner_name("sam3"), Some("sam3"));
542        assert_eq!(arch_runner_name("flux"), Some("flux2"));
543        assert_eq!(arch_runner_name("vjepa2"), Some("vjepa2"));
544        assert_eq!(arch_runner_name("w2v-bert"), Some("wav2vec2-bert"));
545    }
546
547    #[test]
548    fn arch_runner_returns_none_for_embed_and_unknown() {
549        // Embed families aren't in the rlx-run dispatch table today.
550        assert_eq!(arch_runner_name("bert"), None);
551        assert_eq!(arch_runner_name("nomic-bert"), None);
552        assert_eq!(arch_runner_name("totally-fake-arch"), None);
553    }
554
555    #[test]
556    fn known_unimplemented_covers_plan_families() {
557        // M4 — Llama-shaped (real llama.cpp tags)
558        assert_eq!(
559            known_unimplemented_arch("mistral3").map(|u| u.milestone),
560            Some("M4")
561        );
562        assert_eq!(
563            known_unimplemented_arch("phi3").map(|u| u.milestone),
564            Some("M4")
565        );
566        assert_eq!(
567            known_unimplemented_arch("phi4").map(|u| u.milestone),
568            Some("M4")
569        );
570        assert_eq!(
571            known_unimplemented_arch("bonsai").map(|u| u.milestone),
572            Some("M4")
573        );
574        // M5 — MoE / SSM
575        assert_eq!(
576            known_unimplemented_arch("minimax-m2").map(|u| u.milestone),
577            Some("M5")
578        );
579        assert_eq!(
580            known_unimplemented_arch("glm4").map(|u| u.milestone),
581            Some("M5")
582        );
583        assert_eq!(
584            known_unimplemented_arch("nemotron_h").map(|u| u.milestone),
585            Some("M5")
586        );
587        // M6 — MTP
588        assert_eq!(
589            known_unimplemented_arch("qwen3_mtp").map(|u| u.milestone),
590            Some("M6")
591        );
592        // M7 — VL
593        assert_eq!(
594            known_unimplemented_arch("qwen3vl").map(|u| u.milestone),
595            Some("M7")
596        );
597        // Implemented or unknown — plain `mistral` is NOT a llama.cpp arch
598        // tag (Mistral 1/2 use `llama`), so it should not be flagged.
599        assert_eq!(known_unimplemented_arch("qwen3"), None);
600        assert_eq!(known_unimplemented_arch("mistral"), None);
601        assert_eq!(known_unimplemented_arch("totally-fake"), None);
602    }
603
604    #[test]
605    fn auto_sniff_error_points_at_milestone_for_known_unimplemented() {
606        // Build a tiny mistral.gguf and check the error message.
607        let mut buf: Vec<u8> = Vec::new();
608        buf.extend_from_slice(&rlx_gguf::GGUF_MAGIC.to_le_bytes());
609        buf.extend_from_slice(&3u32.to_le_bytes());
610        buf.extend_from_slice(&1u64.to_le_bytes());
611        buf.extend_from_slice(&1u64.to_le_bytes());
612        let k = "general.architecture";
613        buf.extend_from_slice(&(k.len() as u64).to_le_bytes());
614        buf.extend_from_slice(k.as_bytes());
615        buf.extend_from_slice(&8u32.to_le_bytes());
616        let v = "mistral3";
617        buf.extend_from_slice(&(v.len() as u64).to_le_bytes());
618        buf.extend_from_slice(v.as_bytes());
619        let name = "w";
620        buf.extend_from_slice(&(name.len() as u64).to_le_bytes());
621        buf.extend_from_slice(name.as_bytes());
622        buf.extend_from_slice(&1u32.to_le_bytes());
623        buf.extend_from_slice(&4u64.to_le_bytes());
624        buf.extend_from_slice(&(rlx_gguf::GgmlType::F32 as u32).to_le_bytes());
625        buf.extend_from_slice(&0u64.to_le_bytes());
626        while !buf
627            .len()
628            .is_multiple_of(rlx_gguf::DEFAULT_ALIGNMENT as usize)
629        {
630            buf.push(0);
631        }
632        for _ in 0..4 {
633            buf.extend_from_slice(&1.0f32.to_le_bytes());
634        }
635        let path = std::env::temp_dir().join("rlx_auto_dispatch_mistral3_hint.gguf");
636        std::fs::write(&path, &buf).unwrap();
637        let err = auto_sniff(&path).expect_err("should error");
638        let s = format!("{err:#}");
639        assert!(s.contains("Mistral"), "expected family name in error: {s}");
640        assert!(s.contains("M4"), "expected milestone tag in error: {s}");
641        std::fs::remove_file(&path).ok();
642    }
643
644    #[test]
645    fn model_type_runner_maps_known() {
646        assert_eq!(model_type_runner_name("qwen3"), Some("qwen3"));
647        assert_eq!(model_type_runner_name("qwen3_moe"), Some("qwen3"));
648        assert_eq!(model_type_runner_name("llama"), Some("llama32"));
649        assert_eq!(model_type_runner_name("gemma3"), Some("gemma"));
650        // Gemma 4 dense + E2B mobile route to the gemma runner.
651        assert_eq!(model_type_runner_name("gemma4"), Some("gemma"));
652        assert_eq!(model_type_runner_name("gemma4_text"), Some("gemma"));
653        assert_eq!(model_type_runner_name("gemma4_unified"), Some("gemma"));
654        // Gemma 4 dense is no longer flagged unimplemented; MoE A4B still is.
655        assert!(known_unimplemented_arch("gemma4").is_none());
656        assert!(known_unimplemented_arch("gemma4moe").is_some());
657        assert_eq!(
658            model_type_runner_name("dinov2_with_registers"),
659            Some("dinov2")
660        );
661        assert_eq!(model_type_runner_name("whisper"), Some("whisper"));
662        assert_eq!(model_type_runner_name("unknown"), None);
663    }
664
665    /// Builds a minimal GGUF file in a temp dir, then verifies auto_sniff
666    /// picks the right runner name from `general.architecture`.
667    #[test]
668    fn auto_sniff_reads_gguf_arch() {
669        let mut buf: Vec<u8> = Vec::new();
670        buf.extend_from_slice(&rlx_gguf::GGUF_MAGIC.to_le_bytes());
671        buf.extend_from_slice(&3u32.to_le_bytes());
672        buf.extend_from_slice(&1u64.to_le_bytes()); // tensor count
673        buf.extend_from_slice(&1u64.to_le_bytes()); // kv count
674        let write_string = |buf: &mut Vec<u8>, k: &str, v: &str| {
675            buf.extend_from_slice(&(k.len() as u64).to_le_bytes());
676            buf.extend_from_slice(k.as_bytes());
677            buf.extend_from_slice(&8u32.to_le_bytes());
678            buf.extend_from_slice(&(v.len() as u64).to_le_bytes());
679            buf.extend_from_slice(v.as_bytes());
680        };
681        write_string(&mut buf, "general.architecture", "qwen3");
682        // one f32 tensor with 4 elements
683        let name = "w";
684        buf.extend_from_slice(&(name.len() as u64).to_le_bytes());
685        buf.extend_from_slice(name.as_bytes());
686        buf.extend_from_slice(&1u32.to_le_bytes());
687        buf.extend_from_slice(&4u64.to_le_bytes());
688        buf.extend_from_slice(&(rlx_gguf::GgmlType::F32 as u32).to_le_bytes());
689        buf.extend_from_slice(&0u64.to_le_bytes());
690        while !buf
691            .len()
692            .is_multiple_of(rlx_gguf::DEFAULT_ALIGNMENT as usize)
693        {
694            buf.push(0);
695        }
696        for _ in 0..4 {
697            buf.extend_from_slice(&1.0f32.to_le_bytes());
698        }
699        let path = std::env::temp_dir().join("rlx_auto_dispatch_sniff.gguf");
700        std::fs::write(&path, &buf).unwrap();
701        let sniff = auto_sniff(&path).expect("sniff");
702        assert_eq!(sniff.runner_name, "qwen3");
703        match sniff.from {
704            SniffedFrom::GgufArch(a) => assert_eq!(a, "qwen3"),
705            other => panic!("wrong sniff source: {other:?}"),
706        }
707        std::fs::remove_file(&path).ok();
708    }
709
710    /// Register a fake runner under a known name, ask `run_auto` to dispatch
711    /// to it, and capture what argv it received.
712    #[test]
713    fn run_auto_injects_weights_flag_when_missing() {
714        use crate::registry::{ModelRunner, register_runner};
715        use std::sync::{Mutex, OnceLock};
716
717        static CAPTURED: OnceLock<Mutex<Vec<String>>> = OnceLock::new();
718        fn captured() -> &'static Mutex<Vec<String>> {
719            CAPTURED.get_or_init(|| Mutex::new(Vec::new()))
720        }
721
722        struct Capture;
723        impl ModelRunner for Capture {
724            fn name(&self) -> &'static str {
725                "qwen3"
726            }
727            fn description(&self) -> &'static str {
728                "test capture"
729            }
730            fn run(&self, args: &[String]) -> Result<()> {
731                *captured().lock().unwrap() = args.to_vec();
732                Ok(())
733            }
734        }
735        register_runner(Box::new(Capture));
736
737        // Build a minimal qwen3 GGUF in a temp dir.
738        let dir = std::env::temp_dir().join("rlx_auto_dispatch_run_auto");
739        std::fs::create_dir_all(&dir).unwrap();
740        let path = dir.join("model.gguf");
741        let mut buf: Vec<u8> = Vec::new();
742        buf.extend_from_slice(&rlx_gguf::GGUF_MAGIC.to_le_bytes());
743        buf.extend_from_slice(&3u32.to_le_bytes());
744        buf.extend_from_slice(&1u64.to_le_bytes());
745        buf.extend_from_slice(&1u64.to_le_bytes());
746        let k = "general.architecture";
747        buf.extend_from_slice(&(k.len() as u64).to_le_bytes());
748        buf.extend_from_slice(k.as_bytes());
749        buf.extend_from_slice(&8u32.to_le_bytes());
750        let v = "qwen3";
751        buf.extend_from_slice(&(v.len() as u64).to_le_bytes());
752        buf.extend_from_slice(v.as_bytes());
753        let name = "w";
754        buf.extend_from_slice(&(name.len() as u64).to_le_bytes());
755        buf.extend_from_slice(name.as_bytes());
756        buf.extend_from_slice(&1u32.to_le_bytes());
757        buf.extend_from_slice(&4u64.to_le_bytes());
758        buf.extend_from_slice(&(rlx_gguf::GgmlType::F32 as u32).to_le_bytes());
759        buf.extend_from_slice(&0u64.to_le_bytes());
760        while !buf
761            .len()
762            .is_multiple_of(rlx_gguf::DEFAULT_ALIGNMENT as usize)
763        {
764            buf.push(0);
765        }
766        for _ in 0..4 {
767            buf.extend_from_slice(&1.0f32.to_le_bytes());
768        }
769        std::fs::write(&path, &buf).unwrap();
770
771        // Caller passed no --weights → run_auto must inject it.
772        run_auto(&[path.display().to_string(), "--prompt".into(), "hi".into()]).unwrap();
773        let got = captured().lock().unwrap().clone();
774        assert_eq!(
775            got,
776            vec![
777                "--weights".to_string(),
778                path.display().to_string(),
779                "--prompt".into(),
780                "hi".into()
781            ]
782        );
783
784        // Caller already passed --weights → don't inject again.
785        run_auto(&[
786            path.display().to_string(),
787            "--weights".into(),
788            "/other/path".into(),
789            "--prompt".into(),
790            "hi".into(),
791        ])
792        .unwrap();
793        let got = captured().lock().unwrap().clone();
794        assert_eq!(
795            got,
796            vec![
797                "--weights".to_string(),
798                "/other/path".into(),
799                "--prompt".into(),
800                "hi".into(),
801            ]
802        );
803
804        std::fs::remove_dir_all(&dir).ok();
805    }
806
807    #[test]
808    fn auto_sniff_reads_safetensors_sidecar() {
809        let dir = std::env::temp_dir().join("rlx_auto_dispatch_sidecar");
810        std::fs::create_dir_all(&dir).unwrap();
811        let cfg = dir.join("config.json");
812        std::fs::write(&cfg, br#"{"model_type":"llama"}"#).unwrap();
813        let st = dir.join("model.safetensors");
814        // Empty file is fine — sniffer never opens the safetensors payload.
815        std::fs::write(&st, b"").unwrap();
816        let sniff = auto_sniff(&st).expect("sniff");
817        assert_eq!(sniff.runner_name, "llama32");
818        std::fs::remove_dir_all(&dir).ok();
819    }
820}
rlx_cli/auto_dispatch.rs

rlx_cli/
auto_dispatch.rs
