use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::OnceLock;
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
pub(crate) const MODELS_YAML: &str = include_str!("../templates/models.yaml");
pub const MODELS_SCHEMA_VERSION: &str = "1";
pub const OMNI_DEV_MODELS_YAML_ENV: &str = "OMNI_DEV_MODELS_YAML";
const FALLBACK_MAX_OUTPUT_TOKENS: usize = 4096;
const FALLBACK_INPUT_CONTEXT: usize = 100_000;
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Default,
)]
#[serde(rename_all = "lowercase")]
pub enum ModelSource {
#[default]
Embedded,
User,
Project,
Override,
}
impl std::fmt::Display for ModelSource {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(match self {
Self::Embedded => "embedded",
Self::User => "user",
Self::Project => "project",
Self::Override => "override",
})
}
}
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct BetaHeader {
pub key: String,
pub value: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_output_tokens: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input_context: Option<usize>,
}
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct ModelSpec {
pub provider: String,
pub model: String,
pub api_identifier: String,
pub max_output_tokens: usize,
pub input_context: usize,
pub generation: f32,
pub tier: String,
#[serde(default)]
pub legacy: bool,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub beta_headers: Vec<BetaHeader>,
#[serde(default, skip_deserializing)]
pub source: ModelSource,
}
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct TierInfo {
pub description: String,
pub use_cases: Vec<String>,
}
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct DefaultConfig {
pub max_output_tokens: usize,
pub input_context: usize,
}
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct ProviderConfig {
pub name: String,
pub api_base: String,
pub default_model: String,
pub tiers: HashMap<String, TierInfo>,
pub defaults: DefaultConfig,
#[serde(default, skip_deserializing)]
pub source: ModelSource,
}
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct ModelConfiguration {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub version: Option<String>,
pub models: Vec<ModelSpec>,
pub providers: HashMap<String, ProviderConfig>,
}
pub struct ModelRegistry {
config: ModelConfiguration,
by_identifier: HashMap<String, ModelSpec>,
by_provider: HashMap<String, Vec<ModelSpec>>,
}
impl ModelRegistry {
pub fn load() -> Result<Self> {
let override_path = std::env::var(OMNI_DEV_MODELS_YAML_ENV)
.ok()
.filter(|s| !s.is_empty())
.map(PathBuf::from);
let project_path = default_project_path();
let user_path = default_user_path();
Self::load_layered_from_paths(
project_path.as_deref(),
user_path.as_deref(),
override_path.as_deref(),
)
}
pub fn load_layered_from_paths(
project_path: Option<&Path>,
user_path: Option<&Path>,
override_path: Option<&Path>,
) -> Result<Self> {
let mut layers: Vec<(ModelSource, String)> = Vec::new();
layers.push((ModelSource::Embedded, MODELS_YAML.to_string()));
if let Some(path) = override_path {
match read_optional_yaml(path) {
Some(yaml) => layers.push((ModelSource::Override, yaml)),
None => {
tracing::warn!(
"{OMNI_DEV_MODELS_YAML_ENV} points at {} but the file is missing or unreadable; falling back to embedded catalog",
path.display()
);
}
}
} else {
if let Some(path) = user_path {
if let Some(yaml) = read_optional_yaml(path) {
layers.push((ModelSource::User, yaml));
}
}
if let Some(path) = project_path {
if let Some(yaml) = read_optional_yaml(path) {
layers.push((ModelSource::Project, yaml));
}
}
}
Self::from_layers(&layers)
}
pub(crate) fn from_layers(layers: &[(ModelSource, String)]) -> Result<Self> {
let mut merged: serde_yaml::Value =
serde_yaml::Value::Mapping(serde_yaml::Mapping::default());
let mut model_sources: HashMap<String, ModelSource> = HashMap::new();
let mut provider_sources: HashMap<String, ModelSource> = HashMap::new();
let mut declared_versions: Vec<(ModelSource, Option<String>)> = Vec::new();
for (source, yaml) in layers {
let value: serde_yaml::Value = match serde_yaml::from_str(yaml) {
Ok(v) => v,
Err(e) => {
if matches!(source, ModelSource::Embedded) {
return Err(anyhow!(
"Embedded models.yaml is malformed at compile time: {e}"
));
}
tracing::error!(
"Malformed {source} models.yaml: {e}. Falling through to lower-precedence layers."
);
continue;
}
};
let version = value
.get("version")
.and_then(|v| v.as_str())
.map(String::from);
declared_versions.push((*source, version));
merge_layer_into(
&mut merged,
value,
*source,
&mut model_sources,
&mut provider_sources,
);
}
warn_on_version_mismatch(&declared_versions);
let mut config: ModelConfiguration = serde_yaml::from_value(merged)
.map_err(|e| anyhow!("Failed to deserialize merged model configuration: {e}"))?;
for spec in &mut config.models {
spec.source = model_sources
.get(&spec.api_identifier)
.copied()
.unwrap_or_default();
}
for (name, prov) in &mut config.providers {
prov.source = provider_sources.get(name).copied().unwrap_or_default();
}
let mut by_identifier = HashMap::new();
let mut by_provider: HashMap<String, Vec<ModelSpec>> = HashMap::new();
for model in &config.models {
by_identifier.insert(model.api_identifier.clone(), model.clone());
by_provider
.entry(model.provider.clone())
.or_default()
.push(model.clone());
}
Ok(Self {
config,
by_identifier,
by_provider,
})
}
#[must_use]
pub fn config(&self) -> &ModelConfiguration {
&self.config
}
#[must_use]
pub fn get_model_spec(&self, api_identifier: &str) -> Option<&ModelSpec> {
if let Some(spec) = self.by_identifier.get(api_identifier) {
return Some(spec);
}
self.find_model_by_normalized_id(api_identifier)
}
#[must_use]
pub fn get_max_output_tokens(&self, api_identifier: &str) -> usize {
if let Some(spec) = self.get_model_spec(api_identifier) {
return spec.max_output_tokens;
}
if let Some(provider) = self.infer_provider(api_identifier) {
if let Some(provider_config) = self.config.providers.get(&provider) {
return provider_config.defaults.max_output_tokens;
}
}
FALLBACK_MAX_OUTPUT_TOKENS
}
#[must_use]
pub fn get_input_context(&self, api_identifier: &str) -> usize {
if let Some(spec) = self.get_model_spec(api_identifier) {
return spec.input_context;
}
if let Some(provider) = self.infer_provider(api_identifier) {
if let Some(provider_config) = self.config.providers.get(&provider) {
return provider_config.defaults.input_context;
}
}
FALLBACK_INPUT_CONTEXT
}
fn infer_provider(&self, api_identifier: &str) -> Option<String> {
if api_identifier.starts_with("claude") || api_identifier.contains("anthropic") {
Some("claude".to_string())
} else {
None
}
}
fn find_model_by_normalized_id(&self, api_identifier: &str) -> Option<&ModelSpec> {
let core_identifier = self.extract_core_model_identifier(api_identifier);
self.by_identifier.get(&core_identifier)
}
fn extract_core_model_identifier(&self, api_identifier: &str) -> String {
let mut identifier = api_identifier.to_string();
if let Some(dot_pos) = identifier.find('.') {
if identifier[..dot_pos].len() <= 3 {
identifier = identifier[dot_pos + 1..].to_string();
}
}
if identifier.starts_with("anthropic.") {
identifier = identifier["anthropic.".len()..].to_string();
}
if let Some(version_pos) = identifier.rfind("-v") {
if identifier[version_pos..].contains(':') {
identifier = identifier[..version_pos].to_string();
}
}
identifier
}
#[must_use]
pub fn is_legacy_model(&self, api_identifier: &str) -> bool {
self.get_model_spec(api_identifier)
.is_some_and(|spec| spec.legacy)
}
#[must_use]
pub fn get_all_models(&self) -> &[ModelSpec] {
&self.config.models
}
#[must_use]
pub fn get_models_by_provider(&self, provider: &str) -> Vec<&ModelSpec> {
self.by_provider
.get(provider)
.map(|models| models.iter().collect())
.unwrap_or_default()
}
#[must_use]
pub fn get_models_by_provider_and_tier(&self, provider: &str, tier: &str) -> Vec<&ModelSpec> {
self.get_models_by_provider(provider)
.into_iter()
.filter(|model| model.tier == tier)
.collect()
}
#[must_use]
pub fn get_default_model(&self, provider: &str) -> Option<&str> {
self.config
.providers
.get(provider)
.map(|p| p.default_model.as_str())
}
#[must_use]
pub fn get_provider_config(&self, provider: &str) -> Option<&ProviderConfig> {
self.config.providers.get(provider)
}
#[must_use]
pub fn get_tier_info(&self, provider: &str, tier: &str) -> Option<&TierInfo> {
self.config.providers.get(provider)?.tiers.get(tier)
}
#[must_use]
pub fn get_beta_headers(&self, api_identifier: &str) -> &[BetaHeader] {
self.get_model_spec(api_identifier)
.map(|spec| spec.beta_headers.as_slice())
.unwrap_or_default()
}
#[must_use]
pub fn get_max_output_tokens_with_beta(&self, api_identifier: &str, beta_value: &str) -> usize {
if let Some(spec) = self.get_model_spec(api_identifier) {
if let Some(bh) = spec.beta_headers.iter().find(|b| b.value == beta_value) {
if let Some(max) = bh.max_output_tokens {
return max;
}
}
return spec.max_output_tokens;
}
self.get_max_output_tokens(api_identifier)
}
#[must_use]
pub fn get_input_context_with_beta(&self, api_identifier: &str, beta_value: &str) -> usize {
if let Some(spec) = self.get_model_spec(api_identifier) {
if let Some(bh) = spec.beta_headers.iter().find(|b| b.value == beta_value) {
if let Some(ctx) = bh.input_context {
return ctx;
}
}
return spec.input_context;
}
self.get_input_context(api_identifier)
}
}
fn default_project_path() -> Option<PathBuf> {
std::env::current_dir()
.ok()
.map(|cwd| cwd.join(".omni-dev").join("models.yaml"))
}
fn default_user_path() -> Option<PathBuf> {
dirs::home_dir().map(|h| h.join(".omni-dev").join("models.yaml"))
}
fn read_optional_yaml(path: &Path) -> Option<String> {
if !path.exists() {
return None;
}
match std::fs::read_to_string(path) {
Ok(s) => Some(s),
Err(e) => {
tracing::error!(
"Failed to read {}: {e}. Falling through to lower-precedence layers.",
path.display()
);
None
}
}
}
fn merge_layer_into(
dest: &mut serde_yaml::Value,
src: serde_yaml::Value,
source: ModelSource,
model_sources: &mut HashMap<String, ModelSource>,
provider_sources: &mut HashMap<String, ModelSource>,
) {
use serde_yaml::Value;
let Value::Mapping(src_map) = src else {
*dest = src;
return;
};
if !matches!(dest, Value::Mapping(_)) {
*dest = Value::Mapping(serde_yaml::Mapping::new());
}
let Value::Mapping(dest_map) = dest else {
unreachable!("dest is a mapping after the check above");
};
for (k, v) in src_map {
match k.as_str() {
Some("models") => merge_models_into(dest_map, k, v, source, model_sources),
Some("providers") => merge_providers_into(dest_map, k, v, source, provider_sources),
_ => {
dest_map.insert(k, v);
}
}
}
}
fn merge_models_into(
dest_map: &mut serde_yaml::Mapping,
key: serde_yaml::Value,
incoming: serde_yaml::Value,
source: ModelSource,
model_sources: &mut HashMap<String, ModelSource>,
) {
use serde_yaml::Value;
let Value::Sequence(incoming_seq) = incoming else {
dest_map.insert(key, incoming);
return;
};
let dest_value = dest_map
.entry(key)
.or_insert_with(|| Value::Sequence(Vec::new()));
if !matches!(dest_value, Value::Sequence(_)) {
*dest_value = Value::Sequence(Vec::new());
}
let Value::Sequence(dest_seq) = dest_value else {
unreachable!("dest is a sequence after the check above");
};
for entry in incoming_seq {
let api_id = entry
.get("api_identifier")
.and_then(|v| v.as_str())
.map(String::from);
let Some(api_id) = api_id else {
tracing::warn!(
"Skipping model entry without `api_identifier` from {source} models.yaml"
);
continue;
};
if let Some(existing) = dest_seq
.iter_mut()
.find(|e| e.get("api_identifier").and_then(serde_yaml::Value::as_str) == Some(&api_id))
{
deep_merge(existing, entry);
} else {
dest_seq.push(entry);
}
model_sources.insert(api_id, source);
}
}
fn merge_providers_into(
dest_map: &mut serde_yaml::Mapping,
key: serde_yaml::Value,
incoming: serde_yaml::Value,
source: ModelSource,
provider_sources: &mut HashMap<String, ModelSource>,
) {
use serde_yaml::Value;
let Value::Mapping(incoming_providers) = incoming else {
dest_map.insert(key, incoming);
return;
};
let dest_value = dest_map
.entry(key)
.or_insert_with(|| Value::Mapping(serde_yaml::Mapping::new()));
if !matches!(dest_value, Value::Mapping(_)) {
*dest_value = Value::Mapping(serde_yaml::Mapping::new());
}
let Value::Mapping(dest_providers) = dest_value else {
unreachable!("dest is a mapping after the check above");
};
for (pname, pvalue) in incoming_providers {
let pname_str = pname.as_str().map(String::from);
if let Some(existing) = dest_providers.get_mut(&pname) {
deep_merge(existing, pvalue);
} else {
dest_providers.insert(pname.clone(), pvalue);
}
if let Some(name) = pname_str {
provider_sources.insert(name, source);
}
}
}
fn deep_merge(dest: &mut serde_yaml::Value, src: serde_yaml::Value) {
use serde_yaml::Value;
match (dest, src) {
(Value::Mapping(d), Value::Mapping(s)) => {
for (k, v) in s {
if let Some(existing) = d.get_mut(&k) {
deep_merge(existing, v);
} else {
d.insert(k, v);
}
}
}
(d, s) => *d = s,
}
}
fn warn_on_version_mismatch(declared: &[(ModelSource, Option<String>)]) {
for (source, version) in declared {
if matches!(source, ModelSource::Embedded) {
continue;
}
match version {
None => {
tracing::warn!(
"{source} models.yaml has no `version:` field; assuming compatibility with schema version {MODELS_SCHEMA_VERSION}. Add `version: \"{MODELS_SCHEMA_VERSION}\"` to silence this warning."
);
}
Some(v) if v == MODELS_SCHEMA_VERSION => {}
Some(v) => {
tracing::warn!(
"{source} models.yaml declares schema version {v}; this build understands {MODELS_SCHEMA_VERSION}. Continuing — unrecognised fields may be ignored."
);
}
}
}
}
static MODEL_REGISTRY: OnceLock<ModelRegistry> = OnceLock::new();
#[must_use]
pub fn get_model_registry() -> &'static ModelRegistry {
#[allow(clippy::expect_used)] MODEL_REGISTRY.get_or_init(|| ModelRegistry::load().expect("Failed to load model registry"))
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use std::io::Write;
fn embedded_only() -> ModelRegistry {
ModelRegistry::load_layered_from_paths(None, None, None).unwrap()
}
fn write_yaml(dir: &Path, name: &str, contents: &str) -> PathBuf {
let path = dir.join(name);
let mut f = std::fs::File::create(&path).unwrap();
f.write_all(contents.as_bytes()).unwrap();
path
}
#[test]
fn load_model_registry() {
let registry = embedded_only();
assert!(!registry.config.models.is_empty());
assert!(registry.config.providers.contains_key("claude"));
assert_eq!(
registry.config.version.as_deref(),
Some(MODELS_SCHEMA_VERSION)
);
}
#[test]
fn claude_model_lookup() {
let registry = embedded_only();
let opus_spec = registry.get_model_spec("claude-3-opus-20240229");
assert!(opus_spec.is_some());
assert_eq!(opus_spec.unwrap().max_output_tokens, 4096);
assert_eq!(opus_spec.unwrap().provider, "claude");
assert!(registry.is_legacy_model("claude-3-opus-20240229"));
let sonnet45_tokens = registry.get_max_output_tokens("claude-sonnet-4-5-20250929");
assert_eq!(sonnet45_tokens, 64000);
let sonnet4_tokens = registry.get_max_output_tokens("claude-sonnet-4-20250514");
assert_eq!(sonnet4_tokens, 64000);
assert!(registry.is_legacy_model("claude-sonnet-4-20250514"));
let unknown_tokens = registry.get_max_output_tokens("claude-unknown-model");
assert_eq!(unknown_tokens, 4096); }
#[test]
fn unknown_provider_uses_ultimate_fallback() {
let registry = embedded_only();
assert_eq!(
registry.get_max_output_tokens("totally-unknown-vendor-x"),
FALLBACK_MAX_OUTPUT_TOKENS
);
assert_eq!(
registry.get_input_context("totally-unknown-vendor-x"),
FALLBACK_INPUT_CONTEXT
);
}
#[test]
fn provider_filtering() {
let registry = embedded_only();
let claude_models = registry.get_models_by_provider("claude");
assert!(!claude_models.is_empty());
let fast_claude_models = registry.get_models_by_provider_and_tier("claude", "fast");
assert!(!fast_claude_models.is_empty());
let tier_info = registry.get_tier_info("claude", "fast");
assert!(tier_info.is_some());
}
#[test]
fn provider_config() {
let registry = embedded_only();
let claude_config = registry.get_provider_config("claude");
assert!(claude_config.is_some());
assert_eq!(claude_config.unwrap().name, "Anthropic Claude");
}
#[test]
fn default_model_per_provider() {
let registry = embedded_only();
assert_eq!(
registry.get_default_model("claude"),
Some("claude-sonnet-4-6")
);
assert_eq!(registry.get_default_model("openai"), Some("gpt-5-mini"));
assert_eq!(
registry.get_default_model("gemini"),
Some("gemini-2.5-flash")
);
assert_eq!(registry.get_default_model("nonexistent"), None);
}
#[test]
fn normalized_id_matching() {
let registry = embedded_only();
let bedrock_3_7_sonnet = "us.anthropic.claude-3-7-sonnet-20250219-v1:0";
let spec = registry.get_model_spec(bedrock_3_7_sonnet);
assert!(spec.is_some());
assert_eq!(spec.unwrap().api_identifier, "claude-3-7-sonnet-20250219");
assert_eq!(spec.unwrap().max_output_tokens, 64000);
let aws_haiku = "anthropic.claude-3-haiku-20240307-v1:0";
let spec = registry.get_model_spec(aws_haiku);
assert!(spec.is_some());
assert_eq!(spec.unwrap().api_identifier, "claude-3-haiku-20240307");
assert_eq!(spec.unwrap().max_output_tokens, 4096);
let eu_opus = "eu.anthropic.claude-3-opus-20240229-v2:1";
let spec = registry.get_model_spec(eu_opus);
assert!(spec.is_some());
assert_eq!(spec.unwrap().api_identifier, "claude-3-opus-20240229");
assert_eq!(spec.unwrap().max_output_tokens, 4096);
let exact_sonnet45 = "claude-sonnet-4-5-20250929";
let spec = registry.get_model_spec(exact_sonnet45);
assert!(spec.is_some());
assert_eq!(spec.unwrap().max_output_tokens, 64000);
let exact_sonnet4 = "claude-sonnet-4-20250514";
let spec = registry.get_model_spec(exact_sonnet4);
assert!(spec.is_some());
assert_eq!(spec.unwrap().max_output_tokens, 64000);
}
#[test]
fn extract_core_model_identifier() {
let registry = embedded_only();
assert_eq!(
registry.extract_core_model_identifier("us.anthropic.claude-3-7-sonnet-20250219-v1:0"),
"claude-3-7-sonnet-20250219"
);
assert_eq!(
registry.extract_core_model_identifier("anthropic.claude-3-haiku-20240307-v1:0"),
"claude-3-haiku-20240307"
);
assert_eq!(
registry.extract_core_model_identifier("claude-3-opus-20240229"),
"claude-3-opus-20240229"
);
assert_eq!(
registry.extract_core_model_identifier("eu.anthropic.claude-sonnet-4-20250514-v2:1"),
"claude-sonnet-4-20250514"
);
}
#[test]
fn beta_header_lookups() {
let registry = embedded_only();
assert_eq!(registry.get_max_output_tokens("claude-opus-4-6"), 128_000);
assert_eq!(registry.get_input_context("claude-opus-4-6"), 200_000);
assert_eq!(
registry.get_input_context_with_beta("claude-opus-4-6", "context-1m-2025-08-07"),
1_000_000
);
assert_eq!(
registry.get_max_output_tokens_with_beta("claude-opus-4-6", "context-1m-2025-08-07"),
128_000
);
assert_eq!(
registry.get_max_output_tokens_with_beta(
"claude-3-7-sonnet-20250219",
"output-128k-2025-02-19"
),
128_000
);
assert_eq!(
registry.get_max_output_tokens("claude-3-7-sonnet-20250219"),
64000
);
let headers = registry.get_beta_headers("claude-opus-4-6");
assert_eq!(headers.len(), 1);
assert_eq!(headers[0].key, "anthropic-beta");
assert_eq!(headers[0].value, "context-1m-2025-08-07");
let headers = registry.get_beta_headers("claude-3-7-sonnet-20250219");
assert_eq!(headers.len(), 2);
let headers = registry.get_beta_headers("claude-3-haiku-20240307");
assert!(headers.is_empty());
let headers = registry.get_beta_headers("unknown-model");
assert!(headers.is_empty());
}
#[test]
fn beta_lookups_for_unknown_model_fall_through_to_provider_defaults() {
let registry = embedded_only();
assert_eq!(
registry
.get_max_output_tokens_with_beta("claude-unknown-model", "context-1m-2025-08-07"),
4096
);
assert_eq!(
registry.get_input_context_with_beta("claude-unknown-model", "context-1m-2025-08-07"),
200_000
);
}
#[test]
fn embedded_models_default_to_embedded_source() {
let registry = embedded_only();
let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
assert_eq!(spec.source, ModelSource::Embedded);
let provider = registry.get_provider_config("claude").unwrap();
assert_eq!(provider.source, ModelSource::Embedded);
}
#[test]
fn missing_user_and_project_files_fall_through_silently() {
let dir = tempfile::tempdir().unwrap();
let project_path = dir.path().join("missing-project.yaml");
let user_path = dir.path().join("missing-user.yaml");
let registry =
ModelRegistry::load_layered_from_paths(Some(&project_path), Some(&user_path), None)
.unwrap();
let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
assert_eq!(spec.source, ModelSource::Embedded);
assert_eq!(spec.max_output_tokens, 128_000);
}
#[test]
fn user_layer_overrides_embedded_entry() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
version: "1"
models:
- provider: "claude"
model: "Claude Opus 4.6 (custom)"
api_identifier: "claude-opus-4-6"
max_output_tokens: 999999
input_context: 200000
generation: 4.6
tier: "flagship"
"#,
);
let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
assert_eq!(spec.max_output_tokens, 999_999);
assert_eq!(spec.model, "Claude Opus 4.6 (custom)");
assert_eq!(spec.source, ModelSource::User);
}
#[test]
fn project_layer_takes_precedence_over_user_layer() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
version: "1"
models:
- provider: "claude"
model: "From User"
api_identifier: "claude-opus-4-6"
max_output_tokens: 1
input_context: 1
generation: 4.6
tier: "flagship"
"#,
);
let project = write_yaml(
dir.path(),
"project.yaml",
r#"
version: "1"
models:
- provider: "claude"
model: "From Project"
api_identifier: "claude-opus-4-6"
max_output_tokens: 2
input_context: 2
generation: 4.6
tier: "flagship"
"#,
);
let registry =
ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
assert_eq!(spec.model, "From Project");
assert_eq!(spec.max_output_tokens, 2);
assert_eq!(spec.source, ModelSource::Project);
}
#[test]
fn additive_user_entry_is_appended() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
version: "1"
models:
- provider: "claude"
model: "Claude Custom Future"
api_identifier: "claude-future-9000"
max_output_tokens: 250000
input_context: 5000000
generation: 9.0
tier: "flagship"
"#,
);
let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
let spec = registry.get_model_spec("claude-future-9000").unwrap();
assert_eq!(spec.max_output_tokens, 250_000);
assert_eq!(spec.input_context, 5_000_000);
assert_eq!(spec.source, ModelSource::User);
let opus = registry.get_model_spec("claude-opus-4-6").unwrap();
assert_eq!(opus.source, ModelSource::Embedded);
}
#[test]
fn provider_fields_can_be_partially_overridden() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
version: "1"
providers:
claude:
default_model: "claude-opus-4-6"
"#,
);
let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
let claude = registry.get_provider_config("claude").unwrap();
assert_eq!(claude.default_model, "claude-opus-4-6");
assert_eq!(claude.name, "Anthropic Claude");
assert_eq!(claude.api_base, "https://api.anthropic.com/v1");
assert!(claude.tiers.contains_key("flagship"));
assert_eq!(claude.source, ModelSource::User);
}
#[test]
fn malformed_user_yaml_logs_and_falls_through() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
"this: is: definitely: not: valid: yaml: [unbalanced",
);
let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
assert_eq!(spec.source, ModelSource::Embedded);
assert_eq!(spec.max_output_tokens, 128_000);
}
#[test]
fn override_path_short_circuits_user_and_project() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
version: "1"
models:
- provider: "claude"
model: "From User"
api_identifier: "claude-opus-4-6"
max_output_tokens: 1
input_context: 1
generation: 4.6
tier: "flagship"
"#,
);
let project = write_yaml(
dir.path(),
"project.yaml",
r#"
version: "1"
models:
- provider: "claude"
model: "From Project"
api_identifier: "claude-opus-4-6"
max_output_tokens: 2
input_context: 2
generation: 4.6
tier: "flagship"
"#,
);
let override_file = write_yaml(
dir.path(),
"override.yaml",
r#"
version: "1"
models:
- provider: "claude"
model: "From Override"
api_identifier: "claude-opus-4-6"
max_output_tokens: 3
input_context: 3
generation: 4.6
tier: "flagship"
"#,
);
let registry = ModelRegistry::load_layered_from_paths(
Some(&project),
Some(&user),
Some(&override_file),
)
.unwrap();
let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
assert_eq!(spec.model, "From Override");
assert_eq!(spec.max_output_tokens, 3);
assert_eq!(spec.source, ModelSource::Override);
}
#[test]
fn missing_override_path_falls_back_to_embedded() {
let dir = tempfile::tempdir().unwrap();
let missing = dir.path().join("does-not-exist.yaml");
let registry = ModelRegistry::load_layered_from_paths(None, None, Some(&missing)).unwrap();
let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
assert_eq!(spec.source, ModelSource::Embedded);
}
#[test]
fn version_mismatch_is_warned_not_fatal() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
version: "9999"
models:
- provider: "claude"
model: "From Future"
api_identifier: "claude-future-9000"
max_output_tokens: 1
input_context: 1
generation: 9.0
tier: "flagship"
"#,
);
let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
assert!(registry.get_model_spec("claude-future-9000").is_some());
}
#[test]
fn missing_version_is_accepted() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
models:
- provider: "claude"
model: "Versionless"
api_identifier: "claude-versionless"
max_output_tokens: 1
input_context: 1
generation: 1.0
tier: "flagship"
"#,
);
let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
assert!(registry.get_model_spec("claude-versionless").is_some());
}
#[test]
fn model_entry_without_api_identifier_is_skipped() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
version: "1"
models:
- provider: "claude"
model: "No Id"
max_output_tokens: 1
input_context: 1
generation: 1.0
tier: "flagship"
"#,
);
let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
let opus = registry.get_model_spec("claude-opus-4-6").unwrap();
assert_eq!(opus.source, ModelSource::Embedded);
}
#[test]
fn model_source_display() {
assert_eq!(ModelSource::Embedded.to_string(), "embedded");
assert_eq!(ModelSource::User.to_string(), "user");
assert_eq!(ModelSource::Project.to_string(), "project");
assert_eq!(ModelSource::Override.to_string(), "override");
}
#[test]
fn embedded_yaml_must_not_be_malformed() {
let layers = [(ModelSource::Embedded, "::: not yaml :::".to_string())];
let result = ModelRegistry::from_layers(&layers);
assert!(result.is_err());
}
#[test]
fn user_layer_with_scalar_top_level_returns_error() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(dir.path(), "user.yaml", "\"just a string\"\n");
let result = ModelRegistry::load_layered_from_paths(None, Some(&user), None);
assert!(result.is_err());
}
#[test]
fn user_layer_with_non_sequence_models_returns_error() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
version: "1"
models: 42
"#,
);
let result = ModelRegistry::load_layered_from_paths(None, Some(&user), None);
assert!(result.is_err());
}
#[test]
fn user_layer_with_non_mapping_providers_returns_error() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
version: "1"
providers: 42
"#,
);
let result = ModelRegistry::load_layered_from_paths(None, Some(&user), None);
assert!(result.is_err());
}
#[test]
fn deep_merge_inserts_new_keys_into_existing_mapping() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
version: "1"
providers:
claude:
tiers:
experimental:
description: "Experimental tier"
use_cases: ["bleeding edge"]
"#,
);
let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
let claude = registry.get_provider_config("claude").unwrap();
assert!(claude.tiers.contains_key("flagship"));
assert!(claude.tiers.contains_key("balanced"));
assert!(claude.tiers.contains_key("fast"));
let experimental = claude.tiers.get("experimental").unwrap();
assert_eq!(experimental.description, "Experimental tier");
assert_eq!(experimental.use_cases, vec!["bleeding edge".to_string()]);
}
#[test]
#[cfg(unix)]
fn user_path_pointing_at_a_directory_logs_and_falls_through() {
let dir = tempfile::tempdir().unwrap();
let bogus = dir.path().join("models.yaml");
std::fs::create_dir(&bogus).unwrap();
let registry = ModelRegistry::load_layered_from_paths(None, Some(&bogus), None).unwrap();
let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
assert_eq!(spec.source, ModelSource::Embedded);
}
#[test]
#[cfg(unix)]
fn override_path_pointing_at_a_directory_warns_and_falls_through() {
let dir = tempfile::tempdir().unwrap();
let bogus = dir.path().join("override.yaml");
std::fs::create_dir(&bogus).unwrap();
let registry = ModelRegistry::load_layered_from_paths(None, None, Some(&bogus)).unwrap();
let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
assert_eq!(spec.source, ModelSource::Embedded);
}
#[test]
fn project_layer_recovers_after_user_replaces_top_level_with_scalar() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(dir.path(), "user.yaml", "\"junk\"\n");
let project = write_yaml(
dir.path(),
"project.yaml",
r#"
version: "1"
models:
- provider: "claude"
model: "Project Rescue"
api_identifier: "claude-rescue"
max_output_tokens: 1
input_context: 1
generation: 1.0
tier: "flagship"
providers:
custom-provider:
name: "Custom"
api_base: "https://example.invalid"
default_model: "custom-default"
tiers: {}
defaults:
max_output_tokens: 100
input_context: 1000
"#,
);
let registry =
ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
let spec = registry.get_model_spec("claude-rescue").unwrap();
assert_eq!(spec.source, ModelSource::Project);
}
#[test]
fn project_layer_recovers_after_user_replaces_models_with_scalar() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
version: "1"
models: 42
"#,
);
let project = write_yaml(
dir.path(),
"project.yaml",
r#"
version: "1"
models:
- provider: "claude"
model: "Project Rescue"
api_identifier: "claude-rescue"
max_output_tokens: 1
input_context: 1
generation: 1.0
tier: "flagship"
"#,
);
let registry =
ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
let spec = registry.get_model_spec("claude-rescue").unwrap();
assert_eq!(spec.source, ModelSource::Project);
}
#[test]
fn project_layer_recovers_after_user_replaces_providers_with_scalar() {
let dir = tempfile::tempdir().unwrap();
let user = write_yaml(
dir.path(),
"user.yaml",
r#"
version: "1"
providers: 42
"#,
);
let project = write_yaml(
dir.path(),
"project.yaml",
r#"
version: "1"
providers:
custom-provider:
name: "Custom"
api_base: "https://example.invalid"
default_model: "custom-default"
tiers: {}
defaults:
max_output_tokens: 100
input_context: 1000
"#,
);
let registry =
ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
let provider = registry.get_provider_config("custom-provider").unwrap();
assert_eq!(provider.name, "Custom");
assert_eq!(provider.source, ModelSource::Project);
}
#[test]
fn empty_omni_dev_models_yaml_env_var_is_ignored() {
let resolved: Option<PathBuf> = Some(String::new())
.filter(|s| !s.is_empty())
.map(PathBuf::from);
assert!(resolved.is_none());
let resolved: Option<PathBuf> = Some("/some/path".to_string())
.filter(|s| !s.is_empty())
.map(PathBuf::from);
assert_eq!(resolved.as_deref(), Some(Path::new("/some/path")));
}
}