1use std::collections::HashMap;
12use std::path::{Path, PathBuf};
13use std::sync::OnceLock;
14
15use anyhow::{anyhow, Result};
16use serde::{Deserialize, Serialize};
17
18pub(crate) const MODELS_YAML: &str = include_str!("../templates/models.yaml");
20
21pub const MODELS_SCHEMA_VERSION: &str = "1";
27
28pub const OMNI_DEV_MODELS_YAML_ENV: &str = "OMNI_DEV_MODELS_YAML";
31
32const FALLBACK_MAX_OUTPUT_TOKENS: usize = 4096;
34
35const FALLBACK_INPUT_CONTEXT: usize = 100_000;
37
38#[derive(
40 Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Default,
41)]
42#[serde(rename_all = "lowercase")]
43pub enum ModelSource {
44 #[default]
46 Embedded,
47 User,
49 Project,
51 Override,
53}
54
55impl std::fmt::Display for ModelSource {
56 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57 f.write_str(match self {
58 Self::Embedded => "embedded",
59 Self::User => "user",
60 Self::Project => "project",
61 Self::Override => "override",
62 })
63 }
64}
65
66#[derive(Debug, Deserialize, Serialize, Clone)]
68pub struct BetaHeader {
69 pub key: String,
71 pub value: String,
73 #[serde(default, skip_serializing_if = "Option::is_none")]
75 pub max_output_tokens: Option<usize>,
76 #[serde(default, skip_serializing_if = "Option::is_none")]
78 pub input_context: Option<usize>,
79}
80
81#[derive(Debug, Deserialize, Serialize, Clone)]
83pub struct ModelSpec {
84 pub provider: String,
86 pub model: String,
88 pub api_identifier: String,
90 pub max_output_tokens: usize,
92 pub input_context: usize,
94 pub generation: f32,
96 pub tier: String,
98 #[serde(default)]
100 pub legacy: bool,
101 #[serde(default, skip_serializing_if = "Vec::is_empty")]
103 pub beta_headers: Vec<BetaHeader>,
104 #[serde(default, skip_deserializing)]
107 pub source: ModelSource,
108}
109
110#[derive(Debug, Deserialize, Serialize, Clone)]
112pub struct TierInfo {
113 pub description: String,
115 pub use_cases: Vec<String>,
117}
118
119#[derive(Debug, Deserialize, Serialize, Clone)]
121pub struct DefaultConfig {
122 pub max_output_tokens: usize,
124 pub input_context: usize,
126}
127
128#[derive(Debug, Deserialize, Serialize, Clone)]
130pub struct ProviderConfig {
131 pub name: String,
133 pub api_base: String,
135 pub default_model: String,
137 pub tiers: HashMap<String, TierInfo>,
139 pub defaults: DefaultConfig,
141 #[serde(default, skip_deserializing)]
143 pub source: ModelSource,
144}
145
146#[derive(Debug, Deserialize, Serialize, Clone)]
148pub struct ModelConfiguration {
149 #[serde(default, skip_serializing_if = "Option::is_none")]
151 pub version: Option<String>,
152 pub models: Vec<ModelSpec>,
154 pub providers: HashMap<String, ProviderConfig>,
156}
157
158pub struct ModelRegistry {
160 config: ModelConfiguration,
161 by_identifier: HashMap<String, ModelSpec>,
162 by_provider: HashMap<String, Vec<ModelSpec>>,
163}
164
165impl ModelRegistry {
166 pub fn load() -> Result<Self> {
179 let override_path = std::env::var(OMNI_DEV_MODELS_YAML_ENV)
180 .ok()
181 .filter(|s| !s.is_empty())
182 .map(PathBuf::from);
183 let project_path = default_project_path();
184 let user_path = default_user_path();
185 Self::load_layered_from_paths(
186 project_path.as_deref(),
187 user_path.as_deref(),
188 override_path.as_deref(),
189 )
190 }
191
192 pub fn load_layered_from_paths(
196 project_path: Option<&Path>,
197 user_path: Option<&Path>,
198 override_path: Option<&Path>,
199 ) -> Result<Self> {
200 let mut layers: Vec<(ModelSource, String)> = Vec::new();
201 layers.push((ModelSource::Embedded, MODELS_YAML.to_string()));
202
203 if let Some(path) = override_path {
204 match read_optional_yaml(path) {
205 Some(yaml) => layers.push((ModelSource::Override, yaml)),
206 None => {
207 tracing::warn!(
208 "{OMNI_DEV_MODELS_YAML_ENV} points at {} but the file is missing or unreadable; falling back to embedded catalog",
209 path.display()
210 );
211 }
212 }
213 } else {
214 if let Some(path) = user_path {
215 if let Some(yaml) = read_optional_yaml(path) {
216 layers.push((ModelSource::User, yaml));
217 }
218 }
219 if let Some(path) = project_path {
220 if let Some(yaml) = read_optional_yaml(path) {
221 layers.push((ModelSource::Project, yaml));
222 }
223 }
224 }
225
226 Self::from_layers(&layers)
227 }
228
229 pub(crate) fn from_layers(layers: &[(ModelSource, String)]) -> Result<Self> {
235 let mut merged: serde_yaml::Value =
236 serde_yaml::Value::Mapping(serde_yaml::Mapping::default());
237 let mut model_sources: HashMap<String, ModelSource> = HashMap::new();
238 let mut provider_sources: HashMap<String, ModelSource> = HashMap::new();
239 let mut declared_versions: Vec<(ModelSource, Option<String>)> = Vec::new();
240
241 for (source, yaml) in layers {
242 let value: serde_yaml::Value = match serde_yaml::from_str(yaml) {
243 Ok(v) => v,
244 Err(e) => {
245 if matches!(source, ModelSource::Embedded) {
246 return Err(anyhow!(
247 "Embedded models.yaml is malformed at compile time: {e}"
248 ));
249 }
250 tracing::error!(
251 "Malformed {source} models.yaml: {e}. Falling through to lower-precedence layers."
252 );
253 continue;
254 }
255 };
256
257 let version = value
259 .get("version")
260 .and_then(|v| v.as_str())
261 .map(String::from);
262 declared_versions.push((*source, version));
263
264 merge_layer_into(
265 &mut merged,
266 value,
267 *source,
268 &mut model_sources,
269 &mut provider_sources,
270 );
271 }
272
273 warn_on_version_mismatch(&declared_versions);
274
275 let mut config: ModelConfiguration = serde_yaml::from_value(merged)
276 .map_err(|e| anyhow!("Failed to deserialize merged model configuration: {e}"))?;
277
278 for spec in &mut config.models {
279 spec.source = model_sources
280 .get(&spec.api_identifier)
281 .copied()
282 .unwrap_or_default();
283 }
284 for (name, prov) in &mut config.providers {
285 prov.source = provider_sources.get(name).copied().unwrap_or_default();
286 }
287
288 let mut by_identifier = HashMap::new();
289 let mut by_provider: HashMap<String, Vec<ModelSpec>> = HashMap::new();
290 for model in &config.models {
291 by_identifier.insert(model.api_identifier.clone(), model.clone());
292 by_provider
293 .entry(model.provider.clone())
294 .or_default()
295 .push(model.clone());
296 }
297
298 Ok(Self {
299 config,
300 by_identifier,
301 by_provider,
302 })
303 }
304
305 #[must_use]
307 pub fn config(&self) -> &ModelConfiguration {
308 &self.config
309 }
310
311 #[must_use]
313 pub fn get_model_spec(&self, api_identifier: &str) -> Option<&ModelSpec> {
314 if let Some(spec) = self.by_identifier.get(api_identifier) {
316 return Some(spec);
317 }
318
319 self.find_model_by_normalized_id(api_identifier)
321 }
322
323 #[must_use]
325 pub fn get_max_output_tokens(&self, api_identifier: &str) -> usize {
326 if let Some(spec) = self.get_model_spec(api_identifier) {
327 return spec.max_output_tokens;
328 }
329
330 if let Some(provider) = self.infer_provider(api_identifier) {
332 if let Some(provider_config) = self.config.providers.get(&provider) {
333 return provider_config.defaults.max_output_tokens;
334 }
335 }
336
337 FALLBACK_MAX_OUTPUT_TOKENS
339 }
340
341 #[must_use]
343 pub fn get_input_context(&self, api_identifier: &str) -> usize {
344 if let Some(spec) = self.get_model_spec(api_identifier) {
345 return spec.input_context;
346 }
347
348 if let Some(provider) = self.infer_provider(api_identifier) {
350 if let Some(provider_config) = self.config.providers.get(&provider) {
351 return provider_config.defaults.input_context;
352 }
353 }
354
355 FALLBACK_INPUT_CONTEXT
357 }
358
359 fn infer_provider(&self, api_identifier: &str) -> Option<String> {
361 if api_identifier.starts_with("claude") || api_identifier.contains("anthropic") {
362 Some("claude".to_string())
363 } else {
364 None
365 }
366 }
367
368 fn find_model_by_normalized_id(&self, api_identifier: &str) -> Option<&ModelSpec> {
373 let core_identifier = self.extract_core_model_identifier(api_identifier);
374 self.by_identifier.get(&core_identifier)
375 }
376
377 fn extract_core_model_identifier(&self, api_identifier: &str) -> String {
379 let mut identifier = api_identifier.to_string();
380
381 if let Some(dot_pos) = identifier.find('.') {
383 if identifier[..dot_pos].len() <= 3 {
384 identifier = identifier[dot_pos + 1..].to_string();
386 }
387 }
388
389 if identifier.starts_with("anthropic.") {
391 identifier = identifier["anthropic.".len()..].to_string();
392 }
393
394 if let Some(version_pos) = identifier.rfind("-v") {
396 if identifier[version_pos..].contains(':') {
397 identifier = identifier[..version_pos].to_string();
398 }
399 }
400
401 identifier
402 }
403
404 #[must_use]
406 pub fn is_legacy_model(&self, api_identifier: &str) -> bool {
407 self.get_model_spec(api_identifier)
408 .is_some_and(|spec| spec.legacy)
409 }
410
411 #[must_use]
413 pub fn get_all_models(&self) -> &[ModelSpec] {
414 &self.config.models
415 }
416
417 #[must_use]
419 pub fn get_models_by_provider(&self, provider: &str) -> Vec<&ModelSpec> {
420 self.by_provider
421 .get(provider)
422 .map(|models| models.iter().collect())
423 .unwrap_or_default()
424 }
425
426 #[must_use]
428 pub fn get_models_by_provider_and_tier(&self, provider: &str, tier: &str) -> Vec<&ModelSpec> {
429 self.get_models_by_provider(provider)
430 .into_iter()
431 .filter(|model| model.tier == tier)
432 .collect()
433 }
434
435 #[must_use]
437 pub fn get_default_model(&self, provider: &str) -> Option<&str> {
438 self.config
439 .providers
440 .get(provider)
441 .map(|p| p.default_model.as_str())
442 }
443
444 #[must_use]
446 pub fn get_provider_config(&self, provider: &str) -> Option<&ProviderConfig> {
447 self.config.providers.get(provider)
448 }
449
450 #[must_use]
452 pub fn get_tier_info(&self, provider: &str, tier: &str) -> Option<&TierInfo> {
453 self.config.providers.get(provider)?.tiers.get(tier)
454 }
455
456 #[must_use]
458 pub fn get_beta_headers(&self, api_identifier: &str) -> &[BetaHeader] {
459 self.get_model_spec(api_identifier)
460 .map(|spec| spec.beta_headers.as_slice())
461 .unwrap_or_default()
462 }
463
464 #[must_use]
466 pub fn get_max_output_tokens_with_beta(&self, api_identifier: &str, beta_value: &str) -> usize {
467 if let Some(spec) = self.get_model_spec(api_identifier) {
468 if let Some(bh) = spec.beta_headers.iter().find(|b| b.value == beta_value) {
469 if let Some(max) = bh.max_output_tokens {
470 return max;
471 }
472 }
473 return spec.max_output_tokens;
474 }
475 self.get_max_output_tokens(api_identifier)
476 }
477
478 #[must_use]
480 pub fn get_input_context_with_beta(&self, api_identifier: &str, beta_value: &str) -> usize {
481 if let Some(spec) = self.get_model_spec(api_identifier) {
482 if let Some(bh) = spec.beta_headers.iter().find(|b| b.value == beta_value) {
483 if let Some(ctx) = bh.input_context {
484 return ctx;
485 }
486 }
487 return spec.input_context;
488 }
489 self.get_input_context(api_identifier)
490 }
491}
492
493fn default_project_path() -> Option<PathBuf> {
495 std::env::current_dir()
496 .ok()
497 .map(|cwd| cwd.join(".omni-dev").join("models.yaml"))
498}
499
500fn default_user_path() -> Option<PathBuf> {
502 dirs::home_dir().map(|h| h.join(".omni-dev").join("models.yaml"))
503}
504
505fn read_optional_yaml(path: &Path) -> Option<String> {
508 if !path.exists() {
509 return None;
510 }
511 match std::fs::read_to_string(path) {
512 Ok(s) => Some(s),
513 Err(e) => {
514 tracing::error!(
515 "Failed to read {}: {e}. Falling through to lower-precedence layers.",
516 path.display()
517 );
518 None
519 }
520 }
521}
522
523fn merge_layer_into(
534 dest: &mut serde_yaml::Value,
535 src: serde_yaml::Value,
536 source: ModelSource,
537 model_sources: &mut HashMap<String, ModelSource>,
538 provider_sources: &mut HashMap<String, ModelSource>,
539) {
540 use serde_yaml::Value;
541
542 let Value::Mapping(src_map) = src else {
543 *dest = src;
547 return;
548 };
549
550 if !matches!(dest, Value::Mapping(_)) {
551 *dest = Value::Mapping(serde_yaml::Mapping::new());
552 }
553 let Value::Mapping(dest_map) = dest else {
554 unreachable!("dest is a mapping after the check above");
555 };
556
557 for (k, v) in src_map {
558 match k.as_str() {
559 Some("models") => merge_models_into(dest_map, k, v, source, model_sources),
560 Some("providers") => merge_providers_into(dest_map, k, v, source, provider_sources),
561 _ => {
562 dest_map.insert(k, v);
563 }
564 }
565 }
566}
567
568fn merge_models_into(
569 dest_map: &mut serde_yaml::Mapping,
570 key: serde_yaml::Value,
571 incoming: serde_yaml::Value,
572 source: ModelSource,
573 model_sources: &mut HashMap<String, ModelSource>,
574) {
575 use serde_yaml::Value;
576
577 let Value::Sequence(incoming_seq) = incoming else {
578 dest_map.insert(key, incoming);
580 return;
581 };
582
583 let dest_value = dest_map
584 .entry(key)
585 .or_insert_with(|| Value::Sequence(Vec::new()));
586 if !matches!(dest_value, Value::Sequence(_)) {
587 *dest_value = Value::Sequence(Vec::new());
588 }
589 let Value::Sequence(dest_seq) = dest_value else {
590 unreachable!("dest is a sequence after the check above");
591 };
592
593 for entry in incoming_seq {
594 let api_id = entry
595 .get("api_identifier")
596 .and_then(|v| v.as_str())
597 .map(String::from);
598
599 let Some(api_id) = api_id else {
600 tracing::warn!(
601 "Skipping model entry without `api_identifier` from {source} models.yaml"
602 );
603 continue;
604 };
605
606 if let Some(existing) = dest_seq
607 .iter_mut()
608 .find(|e| e.get("api_identifier").and_then(serde_yaml::Value::as_str) == Some(&api_id))
609 {
610 deep_merge(existing, entry);
611 } else {
612 dest_seq.push(entry);
613 }
614
615 model_sources.insert(api_id, source);
616 }
617}
618
619fn merge_providers_into(
620 dest_map: &mut serde_yaml::Mapping,
621 key: serde_yaml::Value,
622 incoming: serde_yaml::Value,
623 source: ModelSource,
624 provider_sources: &mut HashMap<String, ModelSource>,
625) {
626 use serde_yaml::Value;
627
628 let Value::Mapping(incoming_providers) = incoming else {
629 dest_map.insert(key, incoming);
630 return;
631 };
632
633 let dest_value = dest_map
634 .entry(key)
635 .or_insert_with(|| Value::Mapping(serde_yaml::Mapping::new()));
636 if !matches!(dest_value, Value::Mapping(_)) {
637 *dest_value = Value::Mapping(serde_yaml::Mapping::new());
638 }
639 let Value::Mapping(dest_providers) = dest_value else {
640 unreachable!("dest is a mapping after the check above");
641 };
642
643 for (pname, pvalue) in incoming_providers {
644 let pname_str = pname.as_str().map(String::from);
645
646 if let Some(existing) = dest_providers.get_mut(&pname) {
647 deep_merge(existing, pvalue);
648 } else {
649 dest_providers.insert(pname.clone(), pvalue);
650 }
651
652 if let Some(name) = pname_str {
653 provider_sources.insert(name, source);
654 }
655 }
656}
657
658fn deep_merge(dest: &mut serde_yaml::Value, src: serde_yaml::Value) {
661 use serde_yaml::Value;
662 match (dest, src) {
663 (Value::Mapping(d), Value::Mapping(s)) => {
664 for (k, v) in s {
665 if let Some(existing) = d.get_mut(&k) {
666 deep_merge(existing, v);
667 } else {
668 d.insert(k, v);
669 }
670 }
671 }
672 (d, s) => *d = s,
673 }
674}
675
676fn warn_on_version_mismatch(declared: &[(ModelSource, Option<String>)]) {
679 for (source, version) in declared {
680 if matches!(source, ModelSource::Embedded) {
681 continue;
682 }
683 match version {
684 None => {
685 tracing::warn!(
686 "{source} models.yaml has no `version:` field; assuming compatibility with schema version {MODELS_SCHEMA_VERSION}. Add `version: \"{MODELS_SCHEMA_VERSION}\"` to silence this warning."
687 );
688 }
689 Some(v) if v == MODELS_SCHEMA_VERSION => {}
690 Some(v) => {
691 tracing::warn!(
692 "{source} models.yaml declares schema version {v}; this build understands {MODELS_SCHEMA_VERSION}. Continuing — unrecognised fields may be ignored."
693 );
694 }
695 }
696 }
697}
698
699static MODEL_REGISTRY: OnceLock<ModelRegistry> = OnceLock::new();
701
702#[must_use]
704pub fn get_model_registry() -> &'static ModelRegistry {
705 #[allow(clippy::expect_used)] MODEL_REGISTRY.get_or_init(|| ModelRegistry::load().expect("Failed to load model registry"))
707}
708
709#[cfg(test)]
710#[allow(clippy::unwrap_used, clippy::expect_used)]
711mod tests {
712 use super::*;
713 use std::io::Write;
714
715 fn embedded_only() -> ModelRegistry {
716 ModelRegistry::load_layered_from_paths(None, None, None).unwrap()
717 }
718
719 fn write_yaml(dir: &Path, name: &str, contents: &str) -> PathBuf {
720 let path = dir.join(name);
721 let mut f = std::fs::File::create(&path).unwrap();
722 f.write_all(contents.as_bytes()).unwrap();
723 path
724 }
725
726 #[test]
727 fn load_model_registry() {
728 let registry = embedded_only();
729 assert!(!registry.config.models.is_empty());
730 assert!(registry.config.providers.contains_key("claude"));
731 assert_eq!(
732 registry.config.version.as_deref(),
733 Some(MODELS_SCHEMA_VERSION)
734 );
735 }
736
737 #[test]
738 fn claude_model_lookup() {
739 let registry = embedded_only();
740
741 let opus_spec = registry.get_model_spec("claude-3-opus-20240229");
743 assert!(opus_spec.is_some());
744 assert_eq!(opus_spec.unwrap().max_output_tokens, 4096);
745 assert_eq!(opus_spec.unwrap().provider, "claude");
746 assert!(registry.is_legacy_model("claude-3-opus-20240229"));
747
748 let sonnet45_tokens = registry.get_max_output_tokens("claude-sonnet-4-5-20250929");
750 assert_eq!(sonnet45_tokens, 64000);
751
752 let sonnet4_tokens = registry.get_max_output_tokens("claude-sonnet-4-20250514");
754 assert_eq!(sonnet4_tokens, 64000);
755 assert!(registry.is_legacy_model("claude-sonnet-4-20250514"));
756
757 let unknown_tokens = registry.get_max_output_tokens("claude-unknown-model");
759 assert_eq!(unknown_tokens, 4096); }
761
762 #[test]
763 fn unknown_provider_uses_ultimate_fallback() {
764 let registry = embedded_only();
765
766 assert_eq!(
768 registry.get_max_output_tokens("totally-unknown-vendor-x"),
769 FALLBACK_MAX_OUTPUT_TOKENS
770 );
771 assert_eq!(
772 registry.get_input_context("totally-unknown-vendor-x"),
773 FALLBACK_INPUT_CONTEXT
774 );
775 }
776
777 #[test]
778 fn provider_filtering() {
779 let registry = embedded_only();
780
781 let claude_models = registry.get_models_by_provider("claude");
782 assert!(!claude_models.is_empty());
783
784 let fast_claude_models = registry.get_models_by_provider_and_tier("claude", "fast");
785 assert!(!fast_claude_models.is_empty());
786
787 let tier_info = registry.get_tier_info("claude", "fast");
788 assert!(tier_info.is_some());
789 }
790
791 #[test]
792 fn provider_config() {
793 let registry = embedded_only();
794
795 let claude_config = registry.get_provider_config("claude");
796 assert!(claude_config.is_some());
797 assert_eq!(claude_config.unwrap().name, "Anthropic Claude");
798 }
799
800 #[test]
801 fn default_model_per_provider() {
802 let registry = embedded_only();
803
804 assert_eq!(
805 registry.get_default_model("claude"),
806 Some("claude-sonnet-4-6")
807 );
808 assert_eq!(registry.get_default_model("openai"), Some("gpt-5-mini"));
809 assert_eq!(
810 registry.get_default_model("gemini"),
811 Some("gemini-2.5-flash")
812 );
813 assert_eq!(registry.get_default_model("nonexistent"), None);
814 }
815
816 #[test]
817 fn normalized_id_matching() {
818 let registry = embedded_only();
819
820 let bedrock_3_7_sonnet = "us.anthropic.claude-3-7-sonnet-20250219-v1:0";
822 let spec = registry.get_model_spec(bedrock_3_7_sonnet);
823 assert!(spec.is_some());
824 assert_eq!(spec.unwrap().api_identifier, "claude-3-7-sonnet-20250219");
825 assert_eq!(spec.unwrap().max_output_tokens, 64000);
826
827 let aws_haiku = "anthropic.claude-3-haiku-20240307-v1:0";
829 let spec = registry.get_model_spec(aws_haiku);
830 assert!(spec.is_some());
831 assert_eq!(spec.unwrap().api_identifier, "claude-3-haiku-20240307");
832 assert_eq!(spec.unwrap().max_output_tokens, 4096);
833
834 let eu_opus = "eu.anthropic.claude-3-opus-20240229-v2:1";
836 let spec = registry.get_model_spec(eu_opus);
837 assert!(spec.is_some());
838 assert_eq!(spec.unwrap().api_identifier, "claude-3-opus-20240229");
839 assert_eq!(spec.unwrap().max_output_tokens, 4096);
840
841 let exact_sonnet45 = "claude-sonnet-4-5-20250929";
843 let spec = registry.get_model_spec(exact_sonnet45);
844 assert!(spec.is_some());
845 assert_eq!(spec.unwrap().max_output_tokens, 64000);
846
847 let exact_sonnet4 = "claude-sonnet-4-20250514";
849 let spec = registry.get_model_spec(exact_sonnet4);
850 assert!(spec.is_some());
851 assert_eq!(spec.unwrap().max_output_tokens, 64000);
852 }
853
854 #[test]
855 fn extract_core_model_identifier() {
856 let registry = embedded_only();
857
858 assert_eq!(
860 registry.extract_core_model_identifier("us.anthropic.claude-3-7-sonnet-20250219-v1:0"),
861 "claude-3-7-sonnet-20250219"
862 );
863
864 assert_eq!(
865 registry.extract_core_model_identifier("anthropic.claude-3-haiku-20240307-v1:0"),
866 "claude-3-haiku-20240307"
867 );
868
869 assert_eq!(
870 registry.extract_core_model_identifier("claude-3-opus-20240229"),
871 "claude-3-opus-20240229"
872 );
873
874 assert_eq!(
875 registry.extract_core_model_identifier("eu.anthropic.claude-sonnet-4-20250514-v2:1"),
876 "claude-sonnet-4-20250514"
877 );
878 }
879
880 #[test]
881 fn beta_header_lookups() {
882 let registry = embedded_only();
883
884 assert_eq!(registry.get_max_output_tokens("claude-opus-4-6"), 128_000);
886 assert_eq!(registry.get_input_context("claude-opus-4-6"), 200_000);
887
888 assert_eq!(
890 registry.get_input_context_with_beta("claude-opus-4-6", "context-1m-2025-08-07"),
891 1_000_000
892 );
893 assert_eq!(
895 registry.get_max_output_tokens_with_beta("claude-opus-4-6", "context-1m-2025-08-07"),
896 128_000
897 );
898
899 assert_eq!(
901 registry.get_max_output_tokens_with_beta(
902 "claude-3-7-sonnet-20250219",
903 "output-128k-2025-02-19"
904 ),
905 128_000
906 );
907
908 assert_eq!(
910 registry.get_max_output_tokens("claude-3-7-sonnet-20250219"),
911 64000
912 );
913
914 let headers = registry.get_beta_headers("claude-opus-4-6");
916 assert_eq!(headers.len(), 1);
917 assert_eq!(headers[0].key, "anthropic-beta");
918 assert_eq!(headers[0].value, "context-1m-2025-08-07");
919
920 let headers = registry.get_beta_headers("claude-3-7-sonnet-20250219");
922 assert_eq!(headers.len(), 2);
923
924 let headers = registry.get_beta_headers("claude-3-haiku-20240307");
926 assert!(headers.is_empty());
927
928 let headers = registry.get_beta_headers("unknown-model");
930 assert!(headers.is_empty());
931 }
932
933 #[test]
934 fn beta_lookups_for_unknown_model_fall_through_to_provider_defaults() {
935 let registry = embedded_only();
936
937 assert_eq!(
941 registry
942 .get_max_output_tokens_with_beta("claude-unknown-model", "context-1m-2025-08-07"),
943 4096
944 );
945 assert_eq!(
946 registry.get_input_context_with_beta("claude-unknown-model", "context-1m-2025-08-07"),
947 200_000
948 );
949 }
950
951 #[test]
952 fn embedded_models_default_to_embedded_source() {
953 let registry = embedded_only();
954 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
955 assert_eq!(spec.source, ModelSource::Embedded);
956
957 let provider = registry.get_provider_config("claude").unwrap();
958 assert_eq!(provider.source, ModelSource::Embedded);
959 }
960
961 #[test]
962 fn missing_user_and_project_files_fall_through_silently() {
963 let dir = tempfile::tempdir().unwrap();
964 let project_path = dir.path().join("missing-project.yaml");
965 let user_path = dir.path().join("missing-user.yaml");
966 let registry =
967 ModelRegistry::load_layered_from_paths(Some(&project_path), Some(&user_path), None)
968 .unwrap();
969
970 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
972 assert_eq!(spec.source, ModelSource::Embedded);
973 assert_eq!(spec.max_output_tokens, 128_000);
974 }
975
976 #[test]
977 fn user_layer_overrides_embedded_entry() {
978 let dir = tempfile::tempdir().unwrap();
979 let user = write_yaml(
980 dir.path(),
981 "user.yaml",
982 r#"
983version: "1"
984models:
985 - provider: "claude"
986 model: "Claude Opus 4.6 (custom)"
987 api_identifier: "claude-opus-4-6"
988 max_output_tokens: 999999
989 input_context: 200000
990 generation: 4.6
991 tier: "flagship"
992"#,
993 );
994
995 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
996 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
997 assert_eq!(spec.max_output_tokens, 999_999);
998 assert_eq!(spec.model, "Claude Opus 4.6 (custom)");
999 assert_eq!(spec.source, ModelSource::User);
1000 }
1001
1002 #[test]
1003 fn project_layer_takes_precedence_over_user_layer() {
1004 let dir = tempfile::tempdir().unwrap();
1005 let user = write_yaml(
1006 dir.path(),
1007 "user.yaml",
1008 r#"
1009version: "1"
1010models:
1011 - provider: "claude"
1012 model: "From User"
1013 api_identifier: "claude-opus-4-6"
1014 max_output_tokens: 1
1015 input_context: 1
1016 generation: 4.6
1017 tier: "flagship"
1018"#,
1019 );
1020 let project = write_yaml(
1021 dir.path(),
1022 "project.yaml",
1023 r#"
1024version: "1"
1025models:
1026 - provider: "claude"
1027 model: "From Project"
1028 api_identifier: "claude-opus-4-6"
1029 max_output_tokens: 2
1030 input_context: 2
1031 generation: 4.6
1032 tier: "flagship"
1033"#,
1034 );
1035
1036 let registry =
1037 ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
1038 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1039 assert_eq!(spec.model, "From Project");
1040 assert_eq!(spec.max_output_tokens, 2);
1041 assert_eq!(spec.source, ModelSource::Project);
1042 }
1043
1044 #[test]
1045 fn additive_user_entry_is_appended() {
1046 let dir = tempfile::tempdir().unwrap();
1047 let user = write_yaml(
1048 dir.path(),
1049 "user.yaml",
1050 r#"
1051version: "1"
1052models:
1053 - provider: "claude"
1054 model: "Claude Custom Future"
1055 api_identifier: "claude-future-9000"
1056 max_output_tokens: 250000
1057 input_context: 5000000
1058 generation: 9.0
1059 tier: "flagship"
1060"#,
1061 );
1062
1063 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1064 let spec = registry.get_model_spec("claude-future-9000").unwrap();
1065 assert_eq!(spec.max_output_tokens, 250_000);
1066 assert_eq!(spec.input_context, 5_000_000);
1067 assert_eq!(spec.source, ModelSource::User);
1068
1069 let opus = registry.get_model_spec("claude-opus-4-6").unwrap();
1071 assert_eq!(opus.source, ModelSource::Embedded);
1072 }
1073
1074 #[test]
1075 fn provider_fields_can_be_partially_overridden() {
1076 let dir = tempfile::tempdir().unwrap();
1077 let user = write_yaml(
1080 dir.path(),
1081 "user.yaml",
1082 r#"
1083version: "1"
1084providers:
1085 claude:
1086 default_model: "claude-opus-4-6"
1087"#,
1088 );
1089
1090 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1091 let claude = registry.get_provider_config("claude").unwrap();
1092 assert_eq!(claude.default_model, "claude-opus-4-6");
1093 assert_eq!(claude.name, "Anthropic Claude");
1095 assert_eq!(claude.api_base, "https://api.anthropic.com/v1");
1096 assert!(claude.tiers.contains_key("flagship"));
1097 assert_eq!(claude.source, ModelSource::User);
1099 }
1100
1101 #[test]
1102 fn malformed_user_yaml_logs_and_falls_through() {
1103 let dir = tempfile::tempdir().unwrap();
1104 let user = write_yaml(
1105 dir.path(),
1106 "user.yaml",
1107 "this: is: definitely: not: valid: yaml: [unbalanced",
1108 );
1109
1110 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1111 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1113 assert_eq!(spec.source, ModelSource::Embedded);
1114 assert_eq!(spec.max_output_tokens, 128_000);
1115 }
1116
1117 #[test]
1118 fn override_path_short_circuits_user_and_project() {
1119 let dir = tempfile::tempdir().unwrap();
1120 let user = write_yaml(
1121 dir.path(),
1122 "user.yaml",
1123 r#"
1124version: "1"
1125models:
1126 - provider: "claude"
1127 model: "From User"
1128 api_identifier: "claude-opus-4-6"
1129 max_output_tokens: 1
1130 input_context: 1
1131 generation: 4.6
1132 tier: "flagship"
1133"#,
1134 );
1135 let project = write_yaml(
1136 dir.path(),
1137 "project.yaml",
1138 r#"
1139version: "1"
1140models:
1141 - provider: "claude"
1142 model: "From Project"
1143 api_identifier: "claude-opus-4-6"
1144 max_output_tokens: 2
1145 input_context: 2
1146 generation: 4.6
1147 tier: "flagship"
1148"#,
1149 );
1150 let override_file = write_yaml(
1151 dir.path(),
1152 "override.yaml",
1153 r#"
1154version: "1"
1155models:
1156 - provider: "claude"
1157 model: "From Override"
1158 api_identifier: "claude-opus-4-6"
1159 max_output_tokens: 3
1160 input_context: 3
1161 generation: 4.6
1162 tier: "flagship"
1163"#,
1164 );
1165
1166 let registry = ModelRegistry::load_layered_from_paths(
1167 Some(&project),
1168 Some(&user),
1169 Some(&override_file),
1170 )
1171 .unwrap();
1172 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1173 assert_eq!(spec.model, "From Override");
1174 assert_eq!(spec.max_output_tokens, 3);
1175 assert_eq!(spec.source, ModelSource::Override);
1176 }
1177
1178 #[test]
1179 fn missing_override_path_falls_back_to_embedded() {
1180 let dir = tempfile::tempdir().unwrap();
1181 let missing = dir.path().join("does-not-exist.yaml");
1182 let registry = ModelRegistry::load_layered_from_paths(None, None, Some(&missing)).unwrap();
1183 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1184 assert_eq!(spec.source, ModelSource::Embedded);
1185 }
1186
1187 #[test]
1188 fn version_mismatch_is_warned_not_fatal() {
1189 let dir = tempfile::tempdir().unwrap();
1190 let user = write_yaml(
1191 dir.path(),
1192 "user.yaml",
1193 r#"
1194version: "9999"
1195models:
1196 - provider: "claude"
1197 model: "From Future"
1198 api_identifier: "claude-future-9000"
1199 max_output_tokens: 1
1200 input_context: 1
1201 generation: 9.0
1202 tier: "flagship"
1203"#,
1204 );
1205 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1206 assert!(registry.get_model_spec("claude-future-9000").is_some());
1208 }
1209
1210 #[test]
1211 fn missing_version_is_accepted() {
1212 let dir = tempfile::tempdir().unwrap();
1213 let user = write_yaml(
1214 dir.path(),
1215 "user.yaml",
1216 r#"
1217models:
1218 - provider: "claude"
1219 model: "Versionless"
1220 api_identifier: "claude-versionless"
1221 max_output_tokens: 1
1222 input_context: 1
1223 generation: 1.0
1224 tier: "flagship"
1225"#,
1226 );
1227 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1228 assert!(registry.get_model_spec("claude-versionless").is_some());
1229 }
1230
1231 #[test]
1232 fn model_entry_without_api_identifier_is_skipped() {
1233 let dir = tempfile::tempdir().unwrap();
1234 let user = write_yaml(
1235 dir.path(),
1236 "user.yaml",
1237 r#"
1238version: "1"
1239models:
1240 - provider: "claude"
1241 model: "No Id"
1242 max_output_tokens: 1
1243 input_context: 1
1244 generation: 1.0
1245 tier: "flagship"
1246"#,
1247 );
1248 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1249 let opus = registry.get_model_spec("claude-opus-4-6").unwrap();
1251 assert_eq!(opus.source, ModelSource::Embedded);
1252 }
1253
1254 #[test]
1255 fn model_source_display() {
1256 assert_eq!(ModelSource::Embedded.to_string(), "embedded");
1257 assert_eq!(ModelSource::User.to_string(), "user");
1258 assert_eq!(ModelSource::Project.to_string(), "project");
1259 assert_eq!(ModelSource::Override.to_string(), "override");
1260 }
1261
1262 #[test]
1263 fn embedded_yaml_must_not_be_malformed() {
1264 let layers = [(ModelSource::Embedded, "::: not yaml :::".to_string())];
1266 let result = ModelRegistry::from_layers(&layers);
1267 assert!(result.is_err());
1268 }
1269
1270 #[test]
1271 fn user_layer_with_scalar_top_level_returns_error() {
1272 let dir = tempfile::tempdir().unwrap();
1276 let user = write_yaml(dir.path(), "user.yaml", "\"just a string\"\n");
1277 let result = ModelRegistry::load_layered_from_paths(None, Some(&user), None);
1278 assert!(result.is_err());
1279 }
1280
1281 #[test]
1282 fn user_layer_with_non_sequence_models_returns_error() {
1283 let dir = tempfile::tempdir().unwrap();
1287 let user = write_yaml(
1288 dir.path(),
1289 "user.yaml",
1290 r#"
1291version: "1"
1292models: 42
1293"#,
1294 );
1295 let result = ModelRegistry::load_layered_from_paths(None, Some(&user), None);
1296 assert!(result.is_err());
1297 }
1298
1299 #[test]
1300 fn user_layer_with_non_mapping_providers_returns_error() {
1301 let dir = tempfile::tempdir().unwrap();
1304 let user = write_yaml(
1305 dir.path(),
1306 "user.yaml",
1307 r#"
1308version: "1"
1309providers: 42
1310"#,
1311 );
1312 let result = ModelRegistry::load_layered_from_paths(None, Some(&user), None);
1313 assert!(result.is_err());
1314 }
1315
1316 #[test]
1317 fn deep_merge_inserts_new_keys_into_existing_mapping() {
1318 let dir = tempfile::tempdir().unwrap();
1322 let user = write_yaml(
1323 dir.path(),
1324 "user.yaml",
1325 r#"
1326version: "1"
1327providers:
1328 claude:
1329 tiers:
1330 experimental:
1331 description: "Experimental tier"
1332 use_cases: ["bleeding edge"]
1333"#,
1334 );
1335 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1336 let claude = registry.get_provider_config("claude").unwrap();
1337 assert!(claude.tiers.contains_key("flagship"));
1339 assert!(claude.tiers.contains_key("balanced"));
1340 assert!(claude.tiers.contains_key("fast"));
1341 let experimental = claude.tiers.get("experimental").unwrap();
1343 assert_eq!(experimental.description, "Experimental tier");
1344 assert_eq!(experimental.use_cases, vec!["bleeding edge".to_string()]);
1345 }
1346
1347 #[test]
1348 #[cfg(unix)]
1349 fn user_path_pointing_at_a_directory_logs_and_falls_through() {
1350 let dir = tempfile::tempdir().unwrap();
1353 let bogus = dir.path().join("models.yaml");
1354 std::fs::create_dir(&bogus).unwrap();
1355 let registry = ModelRegistry::load_layered_from_paths(None, Some(&bogus), None).unwrap();
1356 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1357 assert_eq!(spec.source, ModelSource::Embedded);
1358 }
1359
1360 #[test]
1361 #[cfg(unix)]
1362 fn override_path_pointing_at_a_directory_warns_and_falls_through() {
1363 let dir = tempfile::tempdir().unwrap();
1364 let bogus = dir.path().join("override.yaml");
1365 std::fs::create_dir(&bogus).unwrap();
1366 let registry = ModelRegistry::load_layered_from_paths(None, None, Some(&bogus)).unwrap();
1367 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1368 assert_eq!(spec.source, ModelSource::Embedded);
1369 }
1370
1371 #[test]
1372 fn project_layer_recovers_after_user_replaces_top_level_with_scalar() {
1373 let dir = tempfile::tempdir().unwrap();
1379 let user = write_yaml(dir.path(), "user.yaml", "\"junk\"\n");
1380 let project = write_yaml(
1381 dir.path(),
1382 "project.yaml",
1383 r#"
1384version: "1"
1385models:
1386 - provider: "claude"
1387 model: "Project Rescue"
1388 api_identifier: "claude-rescue"
1389 max_output_tokens: 1
1390 input_context: 1
1391 generation: 1.0
1392 tier: "flagship"
1393providers:
1394 custom-provider:
1395 name: "Custom"
1396 api_base: "https://example.invalid"
1397 default_model: "custom-default"
1398 tiers: {}
1399 defaults:
1400 max_output_tokens: 100
1401 input_context: 1000
1402"#,
1403 );
1404 let registry =
1405 ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
1406 let spec = registry.get_model_spec("claude-rescue").unwrap();
1408 assert_eq!(spec.source, ModelSource::Project);
1409 }
1410
1411 #[test]
1412 fn project_layer_recovers_after_user_replaces_models_with_scalar() {
1413 let dir = tempfile::tempdir().unwrap();
1417 let user = write_yaml(
1418 dir.path(),
1419 "user.yaml",
1420 r#"
1421version: "1"
1422models: 42
1423"#,
1424 );
1425 let project = write_yaml(
1426 dir.path(),
1427 "project.yaml",
1428 r#"
1429version: "1"
1430models:
1431 - provider: "claude"
1432 model: "Project Rescue"
1433 api_identifier: "claude-rescue"
1434 max_output_tokens: 1
1435 input_context: 1
1436 generation: 1.0
1437 tier: "flagship"
1438"#,
1439 );
1440 let registry =
1441 ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
1442 let spec = registry.get_model_spec("claude-rescue").unwrap();
1443 assert_eq!(spec.source, ModelSource::Project);
1444 }
1445
1446 #[test]
1447 fn project_layer_recovers_after_user_replaces_providers_with_scalar() {
1448 let dir = tempfile::tempdir().unwrap();
1451 let user = write_yaml(
1452 dir.path(),
1453 "user.yaml",
1454 r#"
1455version: "1"
1456providers: 42
1457"#,
1458 );
1459 let project = write_yaml(
1460 dir.path(),
1461 "project.yaml",
1462 r#"
1463version: "1"
1464providers:
1465 custom-provider:
1466 name: "Custom"
1467 api_base: "https://example.invalid"
1468 default_model: "custom-default"
1469 tiers: {}
1470 defaults:
1471 max_output_tokens: 100
1472 input_context: 1000
1473"#,
1474 );
1475 let registry =
1476 ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
1477 let provider = registry.get_provider_config("custom-provider").unwrap();
1478 assert_eq!(provider.name, "Custom");
1479 assert_eq!(provider.source, ModelSource::Project);
1480 }
1481
1482 #[test]
1483 fn empty_omni_dev_models_yaml_env_var_is_ignored() {
1484 let resolved: Option<PathBuf> = Some(String::new())
1488 .filter(|s| !s.is_empty())
1489 .map(PathBuf::from);
1490 assert!(resolved.is_none());
1491 let resolved: Option<PathBuf> = Some("/some/path".to_string())
1492 .filter(|s| !s.is_empty())
1493 .map(PathBuf::from);
1494 assert_eq!(resolved.as_deref(), Some(Path::new("/some/path")));
1495 }
1496}