use std::path::{Path, PathBuf};
use anyhow::{bail, Context, Result};
use crate::serve::auto_pipeline::looks_like_hf_repo_id;
use crate::serve::quant_select::QuantType;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct RepositoryModelSpec {
pub(crate) repository: String,
pub(crate) quant: Option<QuantType>,
pub(crate) selector: Option<String>,
}
impl RepositoryModelSpec {
pub(crate) fn requested_selector(&self) -> Option<&str> {
self.selector.as_deref()
}
pub(crate) fn is_hosted_only(&self) -> bool {
self.selector.is_some() && self.quant.is_none()
}
pub(crate) fn matches_hosted_filename(&self, filename: &str) -> bool {
let Some(selector) = self.requested_selector() else {
return true;
};
if filename == selector {
return true;
}
let Some(basename) = filename.rsplit('/').next() else {
return false;
};
let Some(stem) = basename
.get(..basename.len().saturating_sub(5))
.filter(|_| basename.to_ascii_lowercase().ends_with(".gguf"))
else {
return false;
};
stem.to_ascii_lowercase()
.ends_with(&format!("-{}", selector.to_ascii_lowercase()))
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum ModelSpec {
List,
Path(PathBuf),
Repository(RepositoryModelSpec),
}
pub(crate) fn parse_model_spec(raw: &str) -> Result<ModelSpec> {
if raw == "list" {
return Ok(ModelSpec::List);
}
let path = Path::new(raw);
if path.exists() || is_explicit_path(path) {
return Ok(ModelSpec::Path(path.to_path_buf()));
}
parse_repository_spec(raw).map(ModelSpec::Repository)
}
pub(crate) fn split_repository_quant_suffix(raw: &str) -> (&str, Option<&str>) {
match raw.rsplit_once(':') {
Some((repository, suffix)) if looks_like_hf_repo_id(repository) => {
(repository, Some(suffix))
}
_ => (raw, None),
}
}
pub(crate) fn parse_repository_spec(raw: &str) -> Result<RepositoryModelSpec> {
let (repository, suffix) = split_repository_quant_suffix(raw);
if !looks_like_hf_repo_id(repository) {
bail!(
"model {raw:?} is neither an existing/explicit path nor a Hugging Face repository (expected owner/repository[:SELECTOR])"
);
}
let selector = suffix.map(validate_repository_selector).transpose()?;
let quant = selector
.as_deref()
.and_then(|value| QuantType::from_canonical_str(value).ok());
Ok(RepositoryModelSpec {
repository: repository.to_owned(),
quant,
selector,
})
}
const MAX_REPOSITORY_SELECTOR_BYTES: usize = 512;
fn validate_repository_selector(value: &str) -> Result<String> {
if value.is_empty() || value.len() > MAX_REPOSITORY_SELECTOR_BYTES {
bail!("hosted artifact selector must contain 1..={MAX_REPOSITORY_SELECTOR_BYTES} bytes");
}
if !value.is_ascii()
|| value.starts_with('/')
|| value.ends_with('/')
|| value.contains('\\')
|| value
.split('/')
.any(|part| part.is_empty() || part == "." || part == "..")
|| !value
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b'/'))
{
bail!(
"hosted artifact selector must be a bounded safe label or repository-relative filename"
);
}
Ok(value.to_owned())
}
pub(crate) fn is_explicit_path(path: &Path) -> bool {
if path.is_absolute() {
return true;
}
let rendered = path.as_os_str().to_string_lossy();
matches!(rendered.as_ref(), "." | "..")
|| rendered.starts_with("./")
|| rendered.starts_with("../")
|| rendered.starts_with('~')
}
pub(crate) fn managed_model_root() -> Result<PathBuf> {
managed_model_root_from(
std::env::var_os("XDG_DATA_HOME").map(PathBuf::from),
std::env::var_os("HOME").map(PathBuf::from),
)
}
fn managed_model_root_from(
xdg_data_home: Option<PathBuf>,
home: Option<PathBuf>,
) -> Result<PathBuf> {
let data_home = xdg_data_home
.filter(|path| path.is_absolute())
.or_else(|| {
home.filter(|path| path.is_absolute())
.map(|path| path.join(".local/share"))
})
.context("cannot resolve managed model root: set absolute XDG_DATA_HOME or HOME")?;
Ok(data_home.join("hf2q/models"))
}
pub(crate) fn managed_revision_dir(
root: &Path,
repository: &str,
revision: &str,
) -> Result<PathBuf> {
if revision.len() != 40 || !revision.bytes().all(|byte| byte.is_ascii_hexdigit()) {
bail!("managed model revision must be an exact 40-hex commit, got {revision:?}");
}
if !looks_like_hf_repo_id(repository)
|| repository.contains("..")
|| repository.len() > crate::input::hf_reference::MAX_HF_REPO_ID_BYTES
{
bail!("managed model repository must be a bounded safe owner/repository identity");
}
let (owner, name) = repository
.split_once('/')
.context("managed model repository has no owner")?;
Ok(root.join(owner).join(name).join(revision))
}
pub(crate) fn default_convert_output(
root: &Path,
repository: &str,
revision: &str,
quant_name: &str,
) -> Result<PathBuf> {
let model_name = repository
.rsplit_once('/')
.map(|(_, name)| name)
.filter(|name| !name.is_empty())
.context("repository has no model name")?;
Ok(
managed_revision_dir(root, repository, revision)?.join(format!(
"{model_name}-hf2q-{}.gguf",
quant_name.to_ascii_lowercase()
)),
)
}
pub(crate) fn resolve_output_path(
explicit: Option<&Path>,
default_path: PathBuf,
) -> Result<PathBuf> {
let Some(explicit) = explicit else {
return Ok(default_path);
};
if explicit.exists() && explicit.is_dir() {
let filename = default_path
.file_name()
.context("default model output has no filename")?;
Ok(explicit.join(filename))
} else {
Ok(explicit.to_path_buf())
}
}
#[cfg(test)]
mod tests {
use super::*;
const REVISION: &str = "0123456789abcdef0123456789abcdef01234567";
#[test]
fn repository_quant_suffix_is_case_insensitive_and_exact() {
let parsed = parse_repository_spec("owner/model:Q8_0").unwrap();
assert_eq!(parsed.repository, "owner/model");
assert_eq!(parsed.quant, Some(QuantType::Q8_0));
assert_eq!(parsed.requested_selector(), Some("Q8_0"));
assert_eq!(
parse_repository_spec("owner/model:q4_k_m").unwrap().quant,
Some(QuantType::Q4_K_M)
);
let publisher = parse_repository_spec("owner/model:UD-Q8_K_XL").unwrap();
assert_eq!(publisher.quant, None);
assert_eq!(publisher.requested_selector(), Some("UD-Q8_K_XL"));
assert!(publisher.is_hosted_only());
assert!(publisher.matches_hosted_filename("model-UD-Q8_K_XL.gguf"));
assert!(!publisher.matches_hosted_filename("model-Q8_0.gguf"));
assert!(parse_repository_spec("owner/model:../unsafe.gguf").is_err());
}
#[test]
fn selector_matches_exact_repository_filename_or_literal_stem_suffix() {
let exact = parse_repository_spec("owner/model:gguf/model-Q8_0.gguf").unwrap();
assert!(exact.matches_hosted_filename("gguf/model-Q8_0.gguf"));
assert!(!exact.matches_hosted_filename("other/model-Q8_0.gguf"));
let quant = parse_repository_spec("owner/model:Q8_0").unwrap();
assert!(quant.matches_hosted_filename("model-q8_0.GGUF"));
assert!(quant.matches_hosted_filename("model-MTP-Q8_0.gguf"));
}
#[test]
fn paths_and_list_are_not_reinterpreted_as_repositories() {
assert_eq!(parse_model_spec("list").unwrap(), ModelSpec::List);
assert_eq!(
parse_model_spec("./missing.gguf").unwrap(),
ModelSpec::Path(PathBuf::from("./missing.gguf"))
);
assert!(parse_model_spec("missing.gguf").is_err());
}
#[test]
fn managed_root_honors_xdg_then_home() {
assert_eq!(
managed_model_root_from(Some(PathBuf::from("/data")), Some(PathBuf::from("/home/u")))
.unwrap(),
PathBuf::from("/data/hf2q/models")
);
assert_eq!(
managed_model_root_from(None, Some(PathBuf::from("/home/u"))).unwrap(),
PathBuf::from("/home/u/.local/share/hf2q/models")
);
assert!(managed_model_root_from(None, None).is_err());
}
#[test]
fn default_conversion_path_is_revision_bound() {
let root = tempfile::tempdir().unwrap();
let default =
default_convert_output(root.path(), "owner/My-Model", REVISION, "Q4_K_M").unwrap();
assert_eq!(
default,
root.path().join(format!(
"owner/My-Model/{REVISION}/My-Model-hf2q-q4_k_m.gguf"
))
);
let destination = tempfile::tempdir().unwrap();
assert_eq!(
resolve_output_path(Some(destination.path()), default.clone()).unwrap(),
destination.path().join("My-Model-hf2q-q4_k_m.gguf")
);
assert_eq!(
resolve_output_path(Some(Path::new("/tmp/exact.bin")), default).unwrap(),
PathBuf::from("/tmp/exact.bin")
);
}
#[test]
fn managed_revision_path_is_readable_and_traversal_safe() {
let root = Path::new("/models");
assert_eq!(
managed_revision_dir(root, "jenerallee78/Qwen3.8-27B-Abliterated-SFT", REVISION)
.unwrap(),
root.join("jenerallee78/Qwen3.8-27B-Abliterated-SFT")
.join(REVISION)
);
assert!(managed_revision_dir(root, "owner/..", REVISION).is_err());
assert!(managed_revision_dir(root, "../model", REVISION).is_err());
assert!(managed_revision_dir(root, "owner/model/extra", REVISION).is_err());
}
}