1use std::collections::HashMap;
32use std::path::{Path, PathBuf};
33use std::sync::OnceLock;
34
35use anyhow::{anyhow, Result};
36use serde::{Deserialize, Serialize};
37
38pub(crate) const MODELS_YAML: &str = include_str!("../templates/models.yaml");
40
41pub const MODELS_SCHEMA_VERSION: &str = "1";
47
48pub const OMNI_DEV_MODELS_YAML_ENV: &str = "OMNI_DEV_MODELS_YAML";
51
52const FALLBACK_MAX_OUTPUT_TOKENS: usize = 4096;
54
55const FALLBACK_INPUT_CONTEXT: usize = 100_000;
57
58#[derive(
60 Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Default,
61)]
62#[serde(rename_all = "lowercase")]
63pub enum ModelSource {
64 #[default]
66 Embedded,
67 User,
69 Project,
71 Override,
73}
74
75impl std::fmt::Display for ModelSource {
76 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77 f.write_str(match self {
78 Self::Embedded => "embedded",
79 Self::User => "user",
80 Self::Project => "project",
81 Self::Override => "override",
82 })
83 }
84}
85
86#[derive(Debug, Deserialize, Serialize, Clone)]
100pub struct BetaHeader {
101 pub key: String,
103 pub value: String,
105 #[serde(default, skip_serializing_if = "Option::is_none")]
107 pub max_output_tokens: Option<usize>,
108 #[serde(default, skip_serializing_if = "Option::is_none")]
110 pub input_context: Option<usize>,
111}
112
113#[derive(Debug, Clone, Copy, PartialEq, Eq)]
125pub enum ModelFamily {
126 Fable,
128 Opus,
130 Sonnet,
132 Haiku,
134 Unknown,
136}
137
138impl ModelFamily {
139 #[must_use]
142 pub fn glyph(self) -> &'static str {
143 match self {
144 Self::Fable => "π ",
145 Self::Opus => "π‘",
146 Self::Sonnet => "π’",
147 Self::Haiku => "π΅",
148 Self::Unknown => "βͺ",
149 }
150 }
151
152 #[must_use]
154 pub fn label(self) -> &'static str {
155 match self {
156 Self::Fable => "Fable",
157 Self::Opus => "Opus",
158 Self::Sonnet => "Sonnet",
159 Self::Haiku => "Haiku",
160 Self::Unknown => "Claude",
161 }
162 }
163}
164
165#[derive(Debug, Deserialize, Serialize, Clone)]
192pub struct ModelSpec {
193 pub provider: String,
195 pub model: String,
197 pub api_identifier: String,
199 pub max_output_tokens: usize,
201 pub input_context: usize,
203 pub generation: f32,
205 pub tier: String,
207 #[serde(default)]
209 pub legacy: bool,
210 #[serde(default, skip_serializing_if = "Option::is_none")]
225 pub retires: Option<String>,
226 #[serde(default)]
237 pub supports_structured_output: bool,
238 #[serde(default, skip_serializing_if = "Option::is_none")]
245 pub input_token_price: Option<f64>,
246 #[serde(default, skip_serializing_if = "Option::is_none")]
249 pub output_token_price: Option<f64>,
250 #[serde(default, skip_serializing_if = "Vec::is_empty")]
252 pub beta_headers: Vec<BetaHeader>,
253 #[serde(default, skip_deserializing)]
256 pub source: ModelSource,
257}
258
259#[derive(Debug, Deserialize, Serialize, Clone)]
269pub struct TierInfo {
270 pub description: String,
272 pub use_cases: Vec<String>,
274}
275
276#[derive(Debug, Deserialize, Serialize, Clone)]
287pub struct DefaultConfig {
288 pub max_output_tokens: usize,
290 pub input_context: usize,
292}
293
294#[derive(Debug, Deserialize, Serialize, Clone)]
306pub struct ProviderConfig {
307 pub name: String,
309 pub api_base: String,
311 pub default_model: String,
313 pub tiers: HashMap<String, TierInfo>,
315 pub defaults: DefaultConfig,
317 #[serde(default, skip_deserializing)]
319 pub source: ModelSource,
320}
321
322#[derive(Debug, Deserialize, Serialize, Clone)]
335pub struct ModelConfiguration {
336 #[serde(default, skip_serializing_if = "Option::is_none")]
338 pub version: Option<String>,
339 pub models: Vec<ModelSpec>,
341 pub providers: HashMap<String, ProviderConfig>,
343}
344
345pub struct ModelRegistry {
355 config: ModelConfiguration,
356 by_identifier: HashMap<String, ModelSpec>,
357 by_provider: HashMap<String, Vec<ModelSpec>>,
358}
359
360impl ModelRegistry {
361 pub fn load() -> Result<Self> {
374 let override_path = std::env::var(OMNI_DEV_MODELS_YAML_ENV)
375 .ok()
376 .filter(|s| !s.is_empty())
377 .map(PathBuf::from);
378 let project_path = default_project_path();
379 let user_path = default_user_path();
380 Self::load_layered_from_paths(
381 project_path.as_deref(),
382 user_path.as_deref(),
383 override_path.as_deref(),
384 )
385 }
386
387 pub fn load_layered_from_paths(
391 project_path: Option<&Path>,
392 user_path: Option<&Path>,
393 override_path: Option<&Path>,
394 ) -> Result<Self> {
395 let mut layers: Vec<(ModelSource, String)> = Vec::new();
396 layers.push((ModelSource::Embedded, MODELS_YAML.to_string()));
397
398 if let Some(path) = override_path {
399 match read_optional_yaml(path) {
400 Some(yaml) => layers.push((ModelSource::Override, yaml)),
401 None => {
402 tracing::warn!(
403 "{OMNI_DEV_MODELS_YAML_ENV} points at {} but the file is missing or unreadable; falling back to embedded catalog",
404 path.display()
405 );
406 }
407 }
408 } else {
409 if let Some(path) = user_path {
410 if let Some(yaml) = read_optional_yaml(path) {
411 layers.push((ModelSource::User, yaml));
412 }
413 }
414 if let Some(path) = project_path {
415 if let Some(yaml) = read_optional_yaml(path) {
416 layers.push((ModelSource::Project, yaml));
417 }
418 }
419 }
420
421 Self::from_layers(&layers)
422 }
423
424 pub(crate) fn from_layers(layers: &[(ModelSource, String)]) -> Result<Self> {
430 let mut merged: serde_yaml::Value =
431 serde_yaml::Value::Mapping(serde_yaml::Mapping::default());
432 let mut model_sources: HashMap<String, ModelSource> = HashMap::new();
433 let mut provider_sources: HashMap<String, ModelSource> = HashMap::new();
434 let mut declared_versions: Vec<(ModelSource, Option<String>)> = Vec::new();
435
436 for (source, yaml) in layers {
437 let value: serde_yaml::Value = match serde_yaml::from_str(yaml) {
438 Ok(v) => v,
439 Err(e) => {
440 if matches!(source, ModelSource::Embedded) {
441 return Err(anyhow!(
442 "Embedded models.yaml is malformed at compile time: {e}"
443 ));
444 }
445 tracing::error!(
446 "Malformed {source} models.yaml: {e}. Falling through to lower-precedence layers."
447 );
448 continue;
449 }
450 };
451
452 let version = value
454 .get("version")
455 .and_then(|v| v.as_str())
456 .map(String::from);
457 declared_versions.push((*source, version));
458
459 merge_layer_into(
460 &mut merged,
461 value,
462 *source,
463 &mut model_sources,
464 &mut provider_sources,
465 );
466 }
467
468 warn_on_version_mismatch(&declared_versions);
469
470 let mut config: ModelConfiguration = serde_yaml::from_value(merged)
471 .map_err(|e| anyhow!("Failed to deserialize merged model configuration: {e}"))?;
472
473 for spec in &mut config.models {
474 spec.source = model_sources
475 .get(&spec.api_identifier)
476 .copied()
477 .unwrap_or_default();
478 }
479 for (name, prov) in &mut config.providers {
480 prov.source = provider_sources.get(name).copied().unwrap_or_default();
481 }
482
483 let mut by_identifier = HashMap::new();
484 let mut by_provider: HashMap<String, Vec<ModelSpec>> = HashMap::new();
485 for model in &config.models {
486 by_identifier.insert(model.api_identifier.clone(), model.clone());
487 by_provider
488 .entry(model.provider.clone())
489 .or_default()
490 .push(model.clone());
491 }
492
493 Ok(Self {
494 config,
495 by_identifier,
496 by_provider,
497 })
498 }
499
500 #[must_use]
502 pub fn config(&self) -> &ModelConfiguration {
503 &self.config
504 }
505
506 #[must_use]
508 pub fn get_model_spec(&self, api_identifier: &str) -> Option<&ModelSpec> {
509 if let Some(spec) = self.by_identifier.get(api_identifier) {
511 return Some(spec);
512 }
513
514 self.find_model_by_normalized_id(api_identifier)
516 }
517
518 #[must_use]
520 pub fn get_max_output_tokens(&self, api_identifier: &str) -> usize {
521 if let Some(spec) = self.get_model_spec(api_identifier) {
522 return spec.max_output_tokens;
523 }
524
525 if let Some(provider) = self.infer_provider(api_identifier) {
527 if let Some(provider_config) = self.config.providers.get(&provider) {
528 return provider_config.defaults.max_output_tokens;
529 }
530 }
531
532 FALLBACK_MAX_OUTPUT_TOKENS
534 }
535
536 #[must_use]
538 pub fn get_input_context(&self, api_identifier: &str) -> usize {
539 if let Some(spec) = self.get_model_spec(api_identifier) {
540 return spec.input_context;
541 }
542
543 if let Some(provider) = self.infer_provider(api_identifier) {
545 if let Some(provider_config) = self.config.providers.get(&provider) {
546 return provider_config.defaults.input_context;
547 }
548 }
549
550 FALLBACK_INPUT_CONTEXT
552 }
553
554 #[must_use]
564 pub fn supports_structured_output(&self, api_identifier: &str) -> bool {
565 self.get_model_spec(api_identifier)
566 .is_some_and(|spec| spec.supports_structured_output)
567 }
568
569 fn infer_provider(&self, api_identifier: &str) -> Option<String> {
585 let id = api_identifier;
586 if id.starts_with("claude") || id.contains("anthropic") {
587 Some("claude".to_string())
588 } else if id.starts_with("gpt") || id.starts_with("chatgpt") || is_openai_reasoning_id(id) {
589 Some("openai".to_string())
590 } else if id.starts_with("gemini") {
591 Some("gemini".to_string())
592 } else {
593 None
594 }
595 }
596
597 fn find_model_by_normalized_id(&self, api_identifier: &str) -> Option<&ModelSpec> {
602 let core_identifier = self.extract_core_model_identifier(api_identifier);
603 self.by_identifier.get(&core_identifier)
604 }
605
606 fn extract_core_model_identifier(&self, api_identifier: &str) -> String {
608 let mut identifier = api_identifier.to_string();
609
610 if let Some(dot_pos) = identifier.find('.') {
612 if identifier[..dot_pos].len() <= 3 {
613 identifier = identifier[dot_pos + 1..].to_string();
615 }
616 }
617
618 if identifier.starts_with("anthropic.") {
620 identifier = identifier["anthropic.".len()..].to_string();
621 }
622
623 if let Some(version_pos) = identifier.rfind("-v") {
625 if identifier[version_pos..].contains(':') {
626 identifier = identifier[..version_pos].to_string();
627 }
628 }
629
630 identifier
631 }
632
633 #[must_use]
635 pub fn is_legacy_model(&self, api_identifier: &str) -> bool {
636 self.get_model_spec(api_identifier)
637 .is_some_and(|spec| spec.legacy)
638 }
639
640 #[must_use]
642 pub fn get_all_models(&self) -> &[ModelSpec] {
643 &self.config.models
644 }
645
646 #[must_use]
648 pub fn get_models_by_provider(&self, provider: &str) -> Vec<&ModelSpec> {
649 self.by_provider
650 .get(provider)
651 .map(|models| models.iter().collect())
652 .unwrap_or_default()
653 }
654
655 #[must_use]
661 pub fn is_known_model(&self, identifier: &str) -> bool {
662 self.get_model_spec(identifier).is_some()
663 }
664
665 #[must_use]
671 pub fn known_identifiers(&self, provider: &str) -> Vec<&str> {
672 let mut identifiers: Vec<&str> = self
673 .get_models_by_provider(provider)
674 .into_iter()
675 .filter(|model| !model.legacy)
676 .map(|model| model.api_identifier.as_str())
677 .collect();
678 identifiers.sort_unstable();
679 identifiers
680 }
681
682 #[must_use]
684 pub fn get_models_by_provider_and_tier(&self, provider: &str, tier: &str) -> Vec<&ModelSpec> {
685 self.get_models_by_provider(provider)
686 .into_iter()
687 .filter(|model| model.tier == tier)
688 .collect()
689 }
690
691 #[must_use]
693 pub fn get_default_model(&self, provider: &str) -> Option<&str> {
694 self.config
695 .providers
696 .get(provider)
697 .map(|p| p.default_model.as_str())
698 }
699
700 #[must_use]
702 pub fn get_provider_config(&self, provider: &str) -> Option<&ProviderConfig> {
703 self.config.providers.get(provider)
704 }
705
706 #[must_use]
708 pub fn get_tier_info(&self, provider: &str, tier: &str) -> Option<&TierInfo> {
709 self.config.providers.get(provider)?.tiers.get(tier)
710 }
711
712 #[must_use]
714 pub fn get_beta_headers(&self, api_identifier: &str) -> &[BetaHeader] {
715 self.get_model_spec(api_identifier)
716 .map(|spec| spec.beta_headers.as_slice())
717 .unwrap_or_default()
718 }
719
720 #[must_use]
722 pub fn get_max_output_tokens_with_beta(&self, api_identifier: &str, beta_value: &str) -> usize {
723 if let Some(spec) = self.get_model_spec(api_identifier) {
724 if let Some(bh) = spec.beta_headers.iter().find(|b| b.value == beta_value) {
725 if let Some(max) = bh.max_output_tokens {
726 return max;
727 }
728 }
729 return spec.max_output_tokens;
730 }
731 self.get_max_output_tokens(api_identifier)
732 }
733
734 #[must_use]
736 pub fn get_input_context_with_beta(&self, api_identifier: &str, beta_value: &str) -> usize {
737 if let Some(spec) = self.get_model_spec(api_identifier) {
738 if let Some(bh) = spec.beta_headers.iter().find(|b| b.value == beta_value) {
739 if let Some(ctx) = bh.input_context {
740 return ctx;
741 }
742 }
743 return spec.input_context;
744 }
745 self.get_input_context(api_identifier)
746 }
747
748 #[must_use]
760 pub fn get_model_family(&self, api_identifier: &str) -> ModelFamily {
761 let haystack = self
762 .get_model_spec(api_identifier)
763 .map_or(api_identifier, |spec| spec.model.as_str())
764 .to_ascii_lowercase();
765 if haystack.contains("fable") {
766 ModelFamily::Fable
767 } else if haystack.contains("opus") {
768 ModelFamily::Opus
769 } else if haystack.contains("sonnet") {
770 ModelFamily::Sonnet
771 } else if haystack.contains("haiku") {
772 ModelFamily::Haiku
773 } else {
774 ModelFamily::Unknown
775 }
776 }
777}
778
779fn is_openai_reasoning_id(id: &str) -> bool {
785 let mut chars = id.chars();
786 chars.next() == Some('o') && chars.next().is_some_and(|c| c.is_ascii_digit())
787}
788
789fn default_project_path() -> Option<PathBuf> {
791 std::env::current_dir()
792 .ok()
793 .map(|cwd| cwd.join(".omni-dev").join("models.yaml"))
794}
795
796fn default_user_path() -> Option<PathBuf> {
798 dirs::home_dir().map(|h| h.join(".omni-dev").join("models.yaml"))
799}
800
801fn read_optional_yaml(path: &Path) -> Option<String> {
804 if !path.exists() {
805 return None;
806 }
807 match std::fs::read_to_string(path) {
808 Ok(s) => Some(s),
809 Err(e) => {
810 tracing::error!(
811 "Failed to read {}: {e}. Falling through to lower-precedence layers.",
812 path.display()
813 );
814 None
815 }
816 }
817}
818
819fn merge_layer_into(
830 dest: &mut serde_yaml::Value,
831 src: serde_yaml::Value,
832 source: ModelSource,
833 model_sources: &mut HashMap<String, ModelSource>,
834 provider_sources: &mut HashMap<String, ModelSource>,
835) {
836 use serde_yaml::Value;
837
838 let Value::Mapping(src_map) = src else {
839 *dest = src;
843 return;
844 };
845
846 if !matches!(dest, Value::Mapping(_)) {
847 *dest = Value::Mapping(serde_yaml::Mapping::new());
848 }
849 let Value::Mapping(dest_map) = dest else {
850 unreachable!("dest is a mapping after the check above");
851 };
852
853 for (k, v) in src_map {
854 match k.as_str() {
855 Some("models") => merge_models_into(dest_map, k, v, source, model_sources),
856 Some("providers") => merge_providers_into(dest_map, k, v, source, provider_sources),
857 _ => {
858 dest_map.insert(k, v);
859 }
860 }
861 }
862}
863
864fn merge_models_into(
865 dest_map: &mut serde_yaml::Mapping,
866 key: serde_yaml::Value,
867 incoming: serde_yaml::Value,
868 source: ModelSource,
869 model_sources: &mut HashMap<String, ModelSource>,
870) {
871 use serde_yaml::Value;
872
873 let Value::Sequence(incoming_seq) = incoming else {
874 dest_map.insert(key, incoming);
876 return;
877 };
878
879 let dest_value = dest_map
880 .entry(key)
881 .or_insert_with(|| Value::Sequence(Vec::new()));
882 if !matches!(dest_value, Value::Sequence(_)) {
883 *dest_value = Value::Sequence(Vec::new());
884 }
885 let Value::Sequence(dest_seq) = dest_value else {
886 unreachable!("dest is a sequence after the check above");
887 };
888
889 for entry in incoming_seq {
890 let api_id = entry
891 .get("api_identifier")
892 .and_then(|v| v.as_str())
893 .map(String::from);
894
895 let Some(api_id) = api_id else {
896 tracing::warn!(
897 "Skipping model entry without `api_identifier` from {source} models.yaml"
898 );
899 continue;
900 };
901
902 if let Some(existing) = dest_seq
903 .iter_mut()
904 .find(|e| e.get("api_identifier").and_then(serde_yaml::Value::as_str) == Some(&api_id))
905 {
906 deep_merge(existing, entry);
907 } else {
908 dest_seq.push(entry);
909 }
910
911 model_sources.insert(api_id, source);
912 }
913}
914
915fn merge_providers_into(
916 dest_map: &mut serde_yaml::Mapping,
917 key: serde_yaml::Value,
918 incoming: serde_yaml::Value,
919 source: ModelSource,
920 provider_sources: &mut HashMap<String, ModelSource>,
921) {
922 use serde_yaml::Value;
923
924 let Value::Mapping(incoming_providers) = incoming else {
925 dest_map.insert(key, incoming);
926 return;
927 };
928
929 let dest_value = dest_map
930 .entry(key)
931 .or_insert_with(|| Value::Mapping(serde_yaml::Mapping::new()));
932 if !matches!(dest_value, Value::Mapping(_)) {
933 *dest_value = Value::Mapping(serde_yaml::Mapping::new());
934 }
935 let Value::Mapping(dest_providers) = dest_value else {
936 unreachable!("dest is a mapping after the check above");
937 };
938
939 for (pname, pvalue) in incoming_providers {
940 let pname_str = pname.as_str().map(String::from);
941
942 if let Some(existing) = dest_providers.get_mut(&pname) {
943 deep_merge(existing, pvalue);
944 } else {
945 dest_providers.insert(pname.clone(), pvalue);
946 }
947
948 if let Some(name) = pname_str {
949 provider_sources.insert(name, source);
950 }
951 }
952}
953
954fn deep_merge(dest: &mut serde_yaml::Value, src: serde_yaml::Value) {
957 use serde_yaml::Value;
958 match (dest, src) {
959 (Value::Mapping(d), Value::Mapping(s)) => {
960 for (k, v) in s {
961 if let Some(existing) = d.get_mut(&k) {
962 deep_merge(existing, v);
963 } else {
964 d.insert(k, v);
965 }
966 }
967 }
968 (d, s) => *d = s,
969 }
970}
971
972fn warn_on_version_mismatch(declared: &[(ModelSource, Option<String>)]) {
975 for (source, version) in declared {
976 if matches!(source, ModelSource::Embedded) {
977 continue;
978 }
979 match version {
980 None => {
981 tracing::warn!(
982 "{source} models.yaml has no `version:` field; assuming compatibility with schema version {MODELS_SCHEMA_VERSION}. Add `version: \"{MODELS_SCHEMA_VERSION}\"` to silence this warning."
983 );
984 }
985 Some(v) if v == MODELS_SCHEMA_VERSION => {}
986 Some(v) => {
987 tracing::warn!(
988 "{source} models.yaml declares schema version {v}; this build understands {MODELS_SCHEMA_VERSION}. Continuing β unrecognised fields may be ignored."
989 );
990 }
991 }
992 }
993}
994
995static MODEL_REGISTRY: OnceLock<ModelRegistry> = OnceLock::new();
997
998#[must_use]
1000pub fn get_model_registry() -> &'static ModelRegistry {
1001 #[allow(clippy::expect_used)] MODEL_REGISTRY.get_or_init(|| ModelRegistry::load().expect("Failed to load model registry"))
1003}
1004
1005#[cfg(test)]
1006#[allow(clippy::unwrap_used, clippy::expect_used)]
1007mod tests {
1008 use super::*;
1009 use std::io::Write;
1010
1011 fn embedded_only() -> ModelRegistry {
1012 ModelRegistry::load_layered_from_paths(None, None, None).unwrap()
1013 }
1014
1015 #[test]
1018 fn is_known_model_rejects_unknown_claude_identifier() {
1019 let registry = embedded_only();
1020 assert!(!registry.is_known_model("claude-sonnet-4-8"));
1021 assert!(registry.is_known_model("claude-sonnet-4-6"));
1022 }
1023
1024 #[test]
1025 fn is_known_model_normalizes_bedrock_identifiers() {
1026 let registry = embedded_only();
1027 let bedrock_id = registry
1030 .get_all_models()
1031 .iter()
1032 .find(|m| m.provider == "claude")
1033 .map(|m| format!("us.anthropic.{}", m.api_identifier))
1034 .expect("embedded catalog should list claude models");
1035 assert!(
1036 registry.is_known_model(&bedrock_id),
1037 "{bedrock_id} should normalise to a known model"
1038 );
1039 }
1040
1041 #[test]
1042 fn known_identifiers_lists_non_legacy_claude_models() {
1043 let registry = embedded_only();
1044 let identifiers = registry.known_identifiers("claude");
1045
1046 assert!(!identifiers.is_empty());
1047 assert!(identifiers.contains(&"claude-sonnet-4-6"));
1048 assert!(
1049 identifiers.windows(2).all(|w| w[0] <= w[1]),
1050 "identifiers should be sorted: {identifiers:?}"
1051 );
1052 for id in &identifiers {
1053 let spec = registry
1054 .get_model_spec(id)
1055 .expect("listed model must exist");
1056 assert!(!spec.legacy, "legacy model {id} should not be suggested");
1057 assert_eq!(spec.provider, "claude");
1058 }
1059 }
1060
1061 #[test]
1062 fn known_identifiers_is_empty_for_provider_without_models() {
1063 let registry = embedded_only();
1064 assert!(registry.known_identifiers("ollama").is_empty());
1065 }
1066
1067 #[test]
1070 fn legacy_models_are_known_but_not_suggested() {
1071 let registry = embedded_only();
1072 let legacy = registry
1073 .get_all_models()
1074 .iter()
1075 .find(|m| m.provider == "claude" && m.legacy)
1076 .map(|m| m.api_identifier.clone())
1077 .expect("embedded catalog should list a legacy claude model");
1078
1079 assert!(
1080 registry.is_known_model(&legacy),
1081 "{legacy} must still resolve"
1082 );
1083 assert!(
1084 !registry
1085 .known_identifiers("claude")
1086 .contains(&legacy.as_str()),
1087 "{legacy} should not be suggested"
1088 );
1089 }
1090
1091 fn write_yaml(dir: &Path, name: &str, contents: &str) -> PathBuf {
1092 let path = dir.join(name);
1093 let mut f = std::fs::File::create(&path).unwrap();
1094 f.write_all(contents.as_bytes()).unwrap();
1095 path
1096 }
1097
1098 #[test]
1099 fn load_model_registry() {
1100 let registry = embedded_only();
1101 assert!(!registry.config.models.is_empty());
1102 assert!(registry.config.providers.contains_key("claude"));
1103 assert_eq!(
1104 registry.config.version.as_deref(),
1105 Some(MODELS_SCHEMA_VERSION)
1106 );
1107 }
1108
1109 #[test]
1110 fn claude_model_lookup() {
1111 let registry = embedded_only();
1112
1113 let opus_spec = registry.get_model_spec("claude-3-opus-20240229");
1115 assert!(opus_spec.is_some());
1116 assert_eq!(opus_spec.unwrap().max_output_tokens, 4096);
1117 assert_eq!(opus_spec.unwrap().provider, "claude");
1118 assert!(registry.is_legacy_model("claude-3-opus-20240229"));
1119
1120 let sonnet45_tokens = registry.get_max_output_tokens("claude-sonnet-4-5-20250929");
1122 assert_eq!(sonnet45_tokens, 64000);
1123
1124 let sonnet4_tokens = registry.get_max_output_tokens("claude-sonnet-4-20250514");
1126 assert_eq!(sonnet4_tokens, 64000);
1127 assert!(registry.is_legacy_model("claude-sonnet-4-20250514"));
1128
1129 let unknown_tokens = registry.get_max_output_tokens("claude-unknown-model");
1131 assert_eq!(unknown_tokens, 4096); }
1133
1134 #[test]
1135 fn unknown_provider_uses_ultimate_fallback() {
1136 let registry = embedded_only();
1137
1138 assert_eq!(
1140 registry.get_max_output_tokens("totally-unknown-vendor-x"),
1141 FALLBACK_MAX_OUTPUT_TOKENS
1142 );
1143 assert_eq!(
1144 registry.get_input_context("totally-unknown-vendor-x"),
1145 FALLBACK_INPUT_CONTEXT
1146 );
1147 }
1148
1149 #[test]
1150 fn unknown_openai_model_uses_openai_provider_defaults() {
1151 let registry = embedded_only();
1152
1153 for id in ["gpt-6-ultra", "chatgpt-6-latest", "o5-preview"] {
1158 assert_eq!(
1159 registry.get_max_output_tokens(id),
1160 16384,
1161 "max_output_tokens for {id}"
1162 );
1163 assert_eq!(
1164 registry.get_input_context(id),
1165 128_000,
1166 "input_context for {id}"
1167 );
1168 }
1169 }
1170
1171 #[test]
1172 fn unknown_gemini_model_uses_gemini_provider_defaults() {
1173 let registry = embedded_only();
1174
1175 let id = "gemini-9-pro";
1178 assert_eq!(registry.get_max_output_tokens(id), 8192);
1179 assert_eq!(registry.get_input_context(id), 1_048_576);
1180 }
1181
1182 #[test]
1183 fn provider_filtering() {
1184 let registry = embedded_only();
1185
1186 let claude_models = registry.get_models_by_provider("claude");
1187 assert!(!claude_models.is_empty());
1188
1189 let fast_claude_models = registry.get_models_by_provider_and_tier("claude", "fast");
1190 assert!(!fast_claude_models.is_empty());
1191
1192 let tier_info = registry.get_tier_info("claude", "fast");
1193 assert!(tier_info.is_some());
1194 }
1195
1196 #[test]
1197 fn provider_config() {
1198 let registry = embedded_only();
1199
1200 let claude_config = registry.get_provider_config("claude");
1201 assert!(claude_config.is_some());
1202 assert_eq!(claude_config.unwrap().name, "Anthropic Claude");
1203 }
1204
1205 #[test]
1206 fn default_model_per_provider() {
1207 let registry = embedded_only();
1208
1209 assert_eq!(
1210 registry.get_default_model("claude"),
1211 Some("claude-sonnet-5")
1212 );
1213 assert_eq!(registry.get_default_model("openai"), Some("gpt-5-mini"));
1214 assert_eq!(
1215 registry.get_default_model("gemini"),
1216 Some("gemini-2.5-flash")
1217 );
1218 assert_eq!(registry.get_default_model("nonexistent"), None);
1219 }
1220
1221 #[test]
1222 fn normalized_id_matching() {
1223 let registry = embedded_only();
1224
1225 let bedrock_3_7_sonnet = "us.anthropic.claude-3-7-sonnet-20250219-v1:0";
1227 let spec = registry.get_model_spec(bedrock_3_7_sonnet);
1228 assert!(spec.is_some());
1229 assert_eq!(spec.unwrap().api_identifier, "claude-3-7-sonnet-20250219");
1230 assert_eq!(spec.unwrap().max_output_tokens, 64000);
1231
1232 let aws_haiku = "anthropic.claude-3-haiku-20240307-v1:0";
1234 let spec = registry.get_model_spec(aws_haiku);
1235 assert!(spec.is_some());
1236 assert_eq!(spec.unwrap().api_identifier, "claude-3-haiku-20240307");
1237 assert_eq!(spec.unwrap().max_output_tokens, 4096);
1238
1239 let eu_opus = "eu.anthropic.claude-3-opus-20240229-v2:1";
1241 let spec = registry.get_model_spec(eu_opus);
1242 assert!(spec.is_some());
1243 assert_eq!(spec.unwrap().api_identifier, "claude-3-opus-20240229");
1244 assert_eq!(spec.unwrap().max_output_tokens, 4096);
1245
1246 let exact_sonnet45 = "claude-sonnet-4-5-20250929";
1248 let spec = registry.get_model_spec(exact_sonnet45);
1249 assert!(spec.is_some());
1250 assert_eq!(spec.unwrap().max_output_tokens, 64000);
1251
1252 let exact_sonnet4 = "claude-sonnet-4-20250514";
1254 let spec = registry.get_model_spec(exact_sonnet4);
1255 assert!(spec.is_some());
1256 assert_eq!(spec.unwrap().max_output_tokens, 64000);
1257 }
1258
1259 #[test]
1264 fn supports_structured_output_gates_by_model() {
1265 let registry = embedded_only();
1266
1267 assert!(registry.supports_structured_output("claude-fable-5-1"));
1269 assert!(registry.supports_structured_output("claude-fable-5"));
1270 assert!(registry.supports_structured_output("claude-opus-5-5"));
1271 assert!(registry.supports_structured_output("claude-opus-5"));
1272 assert!(registry.supports_structured_output("claude-opus-4-8"));
1273 assert!(registry.supports_structured_output("claude-opus-4-7"));
1274 assert!(registry.supports_structured_output("claude-sonnet-5"));
1275 assert!(registry.supports_structured_output("claude-sonnet-4-6"));
1276 assert!(registry.supports_structured_output("claude-opus-4-6"));
1277 assert!(registry.supports_structured_output("claude-haiku-4-5-20251001"));
1278 assert!(registry.supports_structured_output("claude-opus-4-5-20251101"));
1279 assert!(
1280 registry.supports_structured_output("us.anthropic.claude-sonnet-4-5-20250929-v1:0"),
1281 "Bedrock/region-prefixed forms must normalize to the flagged spec"
1282 );
1283
1284 assert!(!registry.supports_structured_output("claude-opus-4-1-20250805"));
1286 assert!(!registry.supports_structured_output("claude-3-opus-20240229"));
1287 assert!(!registry.supports_structured_output("claude-sonnet-4-20250514"));
1288
1289 assert!(!registry.supports_structured_output("totally-unknown-model"));
1291 }
1292
1293 #[test]
1297 fn current_generation_models_are_registered() {
1298 let registry = embedded_only();
1299
1300 for id in [
1301 "claude-fable-5-1",
1302 "claude-fable-5",
1303 "claude-opus-5-5",
1304 "claude-opus-5",
1305 "claude-opus-4-8",
1306 "claude-opus-4-7",
1307 "claude-sonnet-5",
1308 ] {
1309 let spec = registry
1310 .get_model_spec(id)
1311 .unwrap_or_else(|| panic!("{id} must be registered"));
1312 assert_eq!(spec.input_context, 1_000_000, "{id} context window");
1313 assert_eq!(spec.max_output_tokens, 128_000, "{id} max output");
1314 assert!(
1315 spec.beta_headers.is_empty(),
1316 "{id} exposes 1M natively and must not rely on a context beta header"
1317 );
1318 assert!(spec.input_token_price.is_some(), "{id} input price");
1319 assert!(spec.output_token_price.is_some(), "{id} output price");
1320 assert!(!spec.legacy, "{id} is current, not legacy");
1321 }
1322
1323 let sonnet5 = registry.get_model_spec("claude-sonnet-5").unwrap();
1326 assert_eq!(sonnet5.input_token_price, Some(3.0));
1327 assert_eq!(sonnet5.output_token_price, Some(15.0));
1328 }
1329
1330 #[test]
1334 fn claude_4_6_limits_are_native_not_beta_gated() {
1335 let registry = embedded_only();
1336
1337 assert_eq!(registry.get_input_context("claude-opus-4-6"), 1_000_000);
1338 assert_eq!(registry.get_input_context("claude-sonnet-4-6"), 1_000_000);
1339 assert_eq!(registry.get_max_output_tokens("claude-sonnet-4-6"), 128_000);
1340 assert_eq!(registry.get_max_output_tokens("claude-opus-4-6"), 128_000);
1341 }
1342
1343 #[test]
1350 fn undated_aliases_resolve_to_real_limits() {
1351 let registry = embedded_only();
1352
1353 for (alias, dated) in [
1354 ("claude-haiku-4-5", "claude-haiku-4-5-20251001"),
1355 ("claude-sonnet-4-5", "claude-sonnet-4-5-20250929"),
1356 ("claude-opus-4-5", "claude-opus-4-5-20251101"),
1357 ] {
1358 assert_eq!(
1359 registry.get_max_output_tokens(alias),
1360 registry.get_max_output_tokens(dated),
1361 "{alias} must resolve to the same output cap as {dated}, \
1362 not the provider fallback"
1363 );
1364 assert_eq!(
1365 registry.get_input_context(alias),
1366 registry.get_input_context(dated),
1367 "{alias} context window must match {dated}"
1368 );
1369 }
1370
1371 assert_eq!(registry.get_max_output_tokens("claude-haiku-4-5"), 64_000);
1373 }
1374
1375 #[test]
1379 fn retires_records_announced_dates_only() {
1380 let registry = embedded_only();
1381
1382 let opus_41 = registry.get_model_spec("claude-opus-4-1-20250805").unwrap();
1383 assert_eq!(opus_41.retires.as_deref(), Some("2026-08-05"));
1384 assert!(opus_41.legacy);
1385
1386 assert_eq!(
1389 registry
1390 .get_model_spec("claude-3-opus-20240229")
1391 .unwrap()
1392 .retires
1393 .as_deref(),
1394 Some("2026-01-05")
1395 );
1396
1397 for id in ["claude-opus-4-20250514", "claude-sonnet-4-20250514"] {
1399 let spec = registry.get_model_spec(id).unwrap();
1400 assert!(spec.legacy, "{id} is deprecated");
1401 assert!(
1402 spec.retires.is_none(),
1403 "{id} has no announced retirement date; absence must not be \
1404 filled in with a guess"
1405 );
1406 }
1407
1408 assert!(registry
1410 .get_model_spec("claude-sonnet-5")
1411 .unwrap()
1412 .retires
1413 .is_none());
1414 }
1415
1416 #[test]
1417 fn extract_core_model_identifier() {
1418 let registry = embedded_only();
1419
1420 assert_eq!(
1422 registry.extract_core_model_identifier("us.anthropic.claude-3-7-sonnet-20250219-v1:0"),
1423 "claude-3-7-sonnet-20250219"
1424 );
1425
1426 assert_eq!(
1427 registry.extract_core_model_identifier("anthropic.claude-3-haiku-20240307-v1:0"),
1428 "claude-3-haiku-20240307"
1429 );
1430
1431 assert_eq!(
1432 registry.extract_core_model_identifier("claude-3-opus-20240229"),
1433 "claude-3-opus-20240229"
1434 );
1435
1436 assert_eq!(
1437 registry.extract_core_model_identifier("eu.anthropic.claude-sonnet-4-20250514-v2:1"),
1438 "claude-sonnet-4-20250514"
1439 );
1440 }
1441
1442 #[test]
1443 fn beta_header_lookups() {
1444 let registry = embedded_only();
1445
1446 assert_eq!(registry.get_max_output_tokens("claude-opus-4-6"), 128_000);
1449 assert_eq!(registry.get_input_context("claude-opus-4-6"), 1_000_000);
1450
1451 assert_eq!(
1456 registry.get_input_context_with_beta("claude-opus-4-6", "context-1m-2025-08-07"),
1457 1_000_000
1458 );
1459 assert_eq!(
1461 registry.get_max_output_tokens_with_beta("claude-opus-4-6", "context-1m-2025-08-07"),
1462 128_000
1463 );
1464
1465 assert_eq!(
1467 registry.get_max_output_tokens_with_beta(
1468 "claude-3-7-sonnet-20250219",
1469 "output-128k-2025-02-19"
1470 ),
1471 128_000
1472 );
1473
1474 assert_eq!(
1476 registry.get_max_output_tokens("claude-3-7-sonnet-20250219"),
1477 64000
1478 );
1479
1480 let headers = registry.get_beta_headers("claude-opus-4-6");
1482 assert_eq!(headers.len(), 1);
1483 assert_eq!(headers[0].key, "anthropic-beta");
1484 assert_eq!(headers[0].value, "context-1m-2025-08-07");
1485
1486 let headers = registry.get_beta_headers("claude-3-7-sonnet-20250219");
1488 assert_eq!(headers.len(), 2);
1489
1490 let headers = registry.get_beta_headers("claude-3-haiku-20240307");
1492 assert!(headers.is_empty());
1493
1494 let headers = registry.get_beta_headers("unknown-model");
1496 assert!(headers.is_empty());
1497 }
1498
1499 #[test]
1500 fn beta_lookups_for_unknown_model_fall_through_to_provider_defaults() {
1501 let registry = embedded_only();
1502
1503 assert_eq!(
1507 registry
1508 .get_max_output_tokens_with_beta("claude-unknown-model", "context-1m-2025-08-07"),
1509 4096
1510 );
1511 assert_eq!(
1512 registry.get_input_context_with_beta("claude-unknown-model", "context-1m-2025-08-07"),
1513 200_000
1514 );
1515 }
1516
1517 #[test]
1518 fn embedded_models_default_to_embedded_source() {
1519 let registry = embedded_only();
1520 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1521 assert_eq!(spec.source, ModelSource::Embedded);
1522
1523 let provider = registry.get_provider_config("claude").unwrap();
1524 assert_eq!(provider.source, ModelSource::Embedded);
1525 }
1526
1527 #[test]
1528 fn missing_user_and_project_files_fall_through_silently() {
1529 let dir = tempfile::tempdir().unwrap();
1530 let project_path = dir.path().join("missing-project.yaml");
1531 let user_path = dir.path().join("missing-user.yaml");
1532 let registry =
1533 ModelRegistry::load_layered_from_paths(Some(&project_path), Some(&user_path), None)
1534 .unwrap();
1535
1536 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1538 assert_eq!(spec.source, ModelSource::Embedded);
1539 assert_eq!(spec.max_output_tokens, 128_000);
1540 }
1541
1542 #[test]
1543 fn user_layer_overrides_embedded_entry() {
1544 let dir = tempfile::tempdir().unwrap();
1545 let user = write_yaml(
1546 dir.path(),
1547 "user.yaml",
1548 r#"
1549version: "1"
1550models:
1551 - provider: "claude"
1552 model: "Claude Opus 4.6 (custom)"
1553 api_identifier: "claude-opus-4-6"
1554 max_output_tokens: 999999
1555 input_context: 200000
1556 generation: 4.6
1557 tier: "flagship"
1558"#,
1559 );
1560
1561 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1562 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1563 assert_eq!(spec.max_output_tokens, 999_999);
1564 assert_eq!(spec.model, "Claude Opus 4.6 (custom)");
1565 assert_eq!(spec.source, ModelSource::User);
1566 }
1567
1568 #[test]
1569 fn project_layer_takes_precedence_over_user_layer() {
1570 let dir = tempfile::tempdir().unwrap();
1571 let user = write_yaml(
1572 dir.path(),
1573 "user.yaml",
1574 r#"
1575version: "1"
1576models:
1577 - provider: "claude"
1578 model: "From User"
1579 api_identifier: "claude-opus-4-6"
1580 max_output_tokens: 1
1581 input_context: 1
1582 generation: 4.6
1583 tier: "flagship"
1584"#,
1585 );
1586 let project = write_yaml(
1587 dir.path(),
1588 "project.yaml",
1589 r#"
1590version: "1"
1591models:
1592 - provider: "claude"
1593 model: "From Project"
1594 api_identifier: "claude-opus-4-6"
1595 max_output_tokens: 2
1596 input_context: 2
1597 generation: 4.6
1598 tier: "flagship"
1599"#,
1600 );
1601
1602 let registry =
1603 ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
1604 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1605 assert_eq!(spec.model, "From Project");
1606 assert_eq!(spec.max_output_tokens, 2);
1607 assert_eq!(spec.source, ModelSource::Project);
1608 }
1609
1610 #[test]
1611 fn additive_user_entry_is_appended() {
1612 let dir = tempfile::tempdir().unwrap();
1613 let user = write_yaml(
1614 dir.path(),
1615 "user.yaml",
1616 r#"
1617version: "1"
1618models:
1619 - provider: "claude"
1620 model: "Claude Custom Future"
1621 api_identifier: "claude-future-9000"
1622 max_output_tokens: 250000
1623 input_context: 5000000
1624 generation: 9.0
1625 tier: "flagship"
1626"#,
1627 );
1628
1629 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1630 let spec = registry.get_model_spec("claude-future-9000").unwrap();
1631 assert_eq!(spec.max_output_tokens, 250_000);
1632 assert_eq!(spec.input_context, 5_000_000);
1633 assert_eq!(spec.source, ModelSource::User);
1634
1635 let opus = registry.get_model_spec("claude-opus-4-6").unwrap();
1637 assert_eq!(opus.source, ModelSource::Embedded);
1638 }
1639
1640 #[test]
1641 fn provider_fields_can_be_partially_overridden() {
1642 let dir = tempfile::tempdir().unwrap();
1643 let user = write_yaml(
1646 dir.path(),
1647 "user.yaml",
1648 r#"
1649version: "1"
1650providers:
1651 claude:
1652 default_model: "claude-opus-4-6"
1653"#,
1654 );
1655
1656 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1657 let claude = registry.get_provider_config("claude").unwrap();
1658 assert_eq!(claude.default_model, "claude-opus-4-6");
1659 assert_eq!(claude.name, "Anthropic Claude");
1661 assert_eq!(claude.api_base, "https://api.anthropic.com/v1");
1662 assert!(claude.tiers.contains_key("flagship"));
1663 assert_eq!(claude.source, ModelSource::User);
1665 }
1666
1667 #[test]
1668 fn malformed_user_yaml_logs_and_falls_through() {
1669 let dir = tempfile::tempdir().unwrap();
1670 let user = write_yaml(
1671 dir.path(),
1672 "user.yaml",
1673 "this: is: definitely: not: valid: yaml: [unbalanced",
1674 );
1675
1676 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1677 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1679 assert_eq!(spec.source, ModelSource::Embedded);
1680 assert_eq!(spec.max_output_tokens, 128_000);
1681 }
1682
1683 #[test]
1684 fn override_path_short_circuits_user_and_project() {
1685 let dir = tempfile::tempdir().unwrap();
1686 let user = write_yaml(
1687 dir.path(),
1688 "user.yaml",
1689 r#"
1690version: "1"
1691models:
1692 - provider: "claude"
1693 model: "From User"
1694 api_identifier: "claude-opus-4-6"
1695 max_output_tokens: 1
1696 input_context: 1
1697 generation: 4.6
1698 tier: "flagship"
1699"#,
1700 );
1701 let project = write_yaml(
1702 dir.path(),
1703 "project.yaml",
1704 r#"
1705version: "1"
1706models:
1707 - provider: "claude"
1708 model: "From Project"
1709 api_identifier: "claude-opus-4-6"
1710 max_output_tokens: 2
1711 input_context: 2
1712 generation: 4.6
1713 tier: "flagship"
1714"#,
1715 );
1716 let override_file = write_yaml(
1717 dir.path(),
1718 "override.yaml",
1719 r#"
1720version: "1"
1721models:
1722 - provider: "claude"
1723 model: "From Override"
1724 api_identifier: "claude-opus-4-6"
1725 max_output_tokens: 3
1726 input_context: 3
1727 generation: 4.6
1728 tier: "flagship"
1729"#,
1730 );
1731
1732 let registry = ModelRegistry::load_layered_from_paths(
1733 Some(&project),
1734 Some(&user),
1735 Some(&override_file),
1736 )
1737 .unwrap();
1738 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1739 assert_eq!(spec.model, "From Override");
1740 assert_eq!(spec.max_output_tokens, 3);
1741 assert_eq!(spec.source, ModelSource::Override);
1742 }
1743
1744 #[test]
1745 fn missing_override_path_falls_back_to_embedded() {
1746 let dir = tempfile::tempdir().unwrap();
1747 let missing = dir.path().join("does-not-exist.yaml");
1748 let registry = ModelRegistry::load_layered_from_paths(None, None, Some(&missing)).unwrap();
1749 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1750 assert_eq!(spec.source, ModelSource::Embedded);
1751 }
1752
1753 #[test]
1754 fn version_mismatch_is_warned_not_fatal() {
1755 let dir = tempfile::tempdir().unwrap();
1756 let user = write_yaml(
1757 dir.path(),
1758 "user.yaml",
1759 r#"
1760version: "9999"
1761models:
1762 - provider: "claude"
1763 model: "From Future"
1764 api_identifier: "claude-future-9000"
1765 max_output_tokens: 1
1766 input_context: 1
1767 generation: 9.0
1768 tier: "flagship"
1769"#,
1770 );
1771 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1772 assert!(registry.get_model_spec("claude-future-9000").is_some());
1774 }
1775
1776 #[test]
1777 fn missing_version_is_accepted() {
1778 let dir = tempfile::tempdir().unwrap();
1779 let user = write_yaml(
1780 dir.path(),
1781 "user.yaml",
1782 r#"
1783models:
1784 - provider: "claude"
1785 model: "Versionless"
1786 api_identifier: "claude-versionless"
1787 max_output_tokens: 1
1788 input_context: 1
1789 generation: 1.0
1790 tier: "flagship"
1791"#,
1792 );
1793 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1794 assert!(registry.get_model_spec("claude-versionless").is_some());
1795 }
1796
1797 #[test]
1798 fn model_entry_without_api_identifier_is_skipped() {
1799 let dir = tempfile::tempdir().unwrap();
1800 let user = write_yaml(
1801 dir.path(),
1802 "user.yaml",
1803 r#"
1804version: "1"
1805models:
1806 - provider: "claude"
1807 model: "No Id"
1808 max_output_tokens: 1
1809 input_context: 1
1810 generation: 1.0
1811 tier: "flagship"
1812"#,
1813 );
1814 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1815 let opus = registry.get_model_spec("claude-opus-4-6").unwrap();
1817 assert_eq!(opus.source, ModelSource::Embedded);
1818 }
1819
1820 #[test]
1821 fn model_source_display() {
1822 assert_eq!(ModelSource::Embedded.to_string(), "embedded");
1823 assert_eq!(ModelSource::User.to_string(), "user");
1824 assert_eq!(ModelSource::Project.to_string(), "project");
1825 assert_eq!(ModelSource::Override.to_string(), "override");
1826 }
1827
1828 #[test]
1829 fn embedded_yaml_must_not_be_malformed() {
1830 let layers = [(ModelSource::Embedded, "::: not yaml :::".to_string())];
1832 let result = ModelRegistry::from_layers(&layers);
1833 assert!(result.is_err());
1834 }
1835
1836 #[test]
1837 fn user_layer_with_scalar_top_level_returns_error() {
1838 let dir = tempfile::tempdir().unwrap();
1842 let user = write_yaml(dir.path(), "user.yaml", "\"just a string\"\n");
1843 let result = ModelRegistry::load_layered_from_paths(None, Some(&user), None);
1844 assert!(result.is_err());
1845 }
1846
1847 #[test]
1848 fn user_layer_with_non_sequence_models_returns_error() {
1849 let dir = tempfile::tempdir().unwrap();
1853 let user = write_yaml(
1854 dir.path(),
1855 "user.yaml",
1856 r#"
1857version: "1"
1858models: 42
1859"#,
1860 );
1861 let result = ModelRegistry::load_layered_from_paths(None, Some(&user), None);
1862 assert!(result.is_err());
1863 }
1864
1865 #[test]
1866 fn user_layer_with_non_mapping_providers_returns_error() {
1867 let dir = tempfile::tempdir().unwrap();
1870 let user = write_yaml(
1871 dir.path(),
1872 "user.yaml",
1873 r#"
1874version: "1"
1875providers: 42
1876"#,
1877 );
1878 let result = ModelRegistry::load_layered_from_paths(None, Some(&user), None);
1879 assert!(result.is_err());
1880 }
1881
1882 #[test]
1883 fn deep_merge_inserts_new_keys_into_existing_mapping() {
1884 let dir = tempfile::tempdir().unwrap();
1888 let user = write_yaml(
1889 dir.path(),
1890 "user.yaml",
1891 r#"
1892version: "1"
1893providers:
1894 claude:
1895 tiers:
1896 experimental:
1897 description: "Experimental tier"
1898 use_cases: ["bleeding edge"]
1899"#,
1900 );
1901 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1902 let claude = registry.get_provider_config("claude").unwrap();
1903 assert!(claude.tiers.contains_key("flagship"));
1905 assert!(claude.tiers.contains_key("balanced"));
1906 assert!(claude.tiers.contains_key("fast"));
1907 let experimental = claude.tiers.get("experimental").unwrap();
1909 assert_eq!(experimental.description, "Experimental tier");
1910 assert_eq!(experimental.use_cases, vec!["bleeding edge".to_string()]);
1911 }
1912
1913 #[test]
1914 #[cfg(unix)]
1915 fn user_path_pointing_at_a_directory_logs_and_falls_through() {
1916 let dir = tempfile::tempdir().unwrap();
1919 let bogus = dir.path().join("models.yaml");
1920 std::fs::create_dir(&bogus).unwrap();
1921 let registry = ModelRegistry::load_layered_from_paths(None, Some(&bogus), None).unwrap();
1922 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1923 assert_eq!(spec.source, ModelSource::Embedded);
1924 }
1925
1926 #[test]
1927 #[cfg(unix)]
1928 fn override_path_pointing_at_a_directory_warns_and_falls_through() {
1929 let dir = tempfile::tempdir().unwrap();
1930 let bogus = dir.path().join("override.yaml");
1931 std::fs::create_dir(&bogus).unwrap();
1932 let registry = ModelRegistry::load_layered_from_paths(None, None, Some(&bogus)).unwrap();
1933 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1934 assert_eq!(spec.source, ModelSource::Embedded);
1935 }
1936
1937 #[test]
1938 fn project_layer_recovers_after_user_replaces_top_level_with_scalar() {
1939 let dir = tempfile::tempdir().unwrap();
1945 let user = write_yaml(dir.path(), "user.yaml", "\"junk\"\n");
1946 let project = write_yaml(
1947 dir.path(),
1948 "project.yaml",
1949 r#"
1950version: "1"
1951models:
1952 - provider: "claude"
1953 model: "Project Rescue"
1954 api_identifier: "claude-rescue"
1955 max_output_tokens: 1
1956 input_context: 1
1957 generation: 1.0
1958 tier: "flagship"
1959providers:
1960 custom-provider:
1961 name: "Custom"
1962 api_base: "https://example.invalid"
1963 default_model: "custom-default"
1964 tiers: {}
1965 defaults:
1966 max_output_tokens: 100
1967 input_context: 1000
1968"#,
1969 );
1970 let registry =
1971 ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
1972 let spec = registry.get_model_spec("claude-rescue").unwrap();
1974 assert_eq!(spec.source, ModelSource::Project);
1975 }
1976
1977 #[test]
1978 fn project_layer_recovers_after_user_replaces_models_with_scalar() {
1979 let dir = tempfile::tempdir().unwrap();
1983 let user = write_yaml(
1984 dir.path(),
1985 "user.yaml",
1986 r#"
1987version: "1"
1988models: 42
1989"#,
1990 );
1991 let project = write_yaml(
1992 dir.path(),
1993 "project.yaml",
1994 r#"
1995version: "1"
1996models:
1997 - provider: "claude"
1998 model: "Project Rescue"
1999 api_identifier: "claude-rescue"
2000 max_output_tokens: 1
2001 input_context: 1
2002 generation: 1.0
2003 tier: "flagship"
2004"#,
2005 );
2006 let registry =
2007 ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
2008 let spec = registry.get_model_spec("claude-rescue").unwrap();
2009 assert_eq!(spec.source, ModelSource::Project);
2010 }
2011
2012 #[test]
2013 fn project_layer_recovers_after_user_replaces_providers_with_scalar() {
2014 let dir = tempfile::tempdir().unwrap();
2017 let user = write_yaml(
2018 dir.path(),
2019 "user.yaml",
2020 r#"
2021version: "1"
2022providers: 42
2023"#,
2024 );
2025 let project = write_yaml(
2026 dir.path(),
2027 "project.yaml",
2028 r#"
2029version: "1"
2030providers:
2031 custom-provider:
2032 name: "Custom"
2033 api_base: "https://example.invalid"
2034 default_model: "custom-default"
2035 tiers: {}
2036 defaults:
2037 max_output_tokens: 100
2038 input_context: 1000
2039"#,
2040 );
2041 let registry =
2042 ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
2043 let provider = registry.get_provider_config("custom-provider").unwrap();
2044 assert_eq!(provider.name, "Custom");
2045 assert_eq!(provider.source, ModelSource::Project);
2046 }
2047
2048 #[test]
2049 fn empty_omni_dev_models_yaml_env_var_is_ignored() {
2050 let resolved: Option<PathBuf> = Some(String::new())
2054 .filter(|s| !s.is_empty())
2055 .map(PathBuf::from);
2056 assert!(resolved.is_none());
2057 let resolved: Option<PathBuf> = Some("/some/path".to_string())
2058 .filter(|s| !s.is_empty())
2059 .map(PathBuf::from);
2060 assert_eq!(resolved.as_deref(), Some(Path::new("/some/path")));
2061 }
2062
2063 #[test]
2064 fn model_family_classifies_registered_claude_identifiers() {
2065 let registry = embedded_only();
2066 assert_eq!(
2067 registry.get_model_family("claude-fable-5"),
2068 ModelFamily::Fable
2069 );
2070 assert_eq!(
2071 registry.get_model_family("claude-opus-4-8"),
2072 ModelFamily::Opus
2073 );
2074 assert_eq!(
2075 registry.get_model_family("claude-sonnet-5"),
2076 ModelFamily::Sonnet
2077 );
2078 assert_eq!(
2079 registry.get_model_family("claude-haiku-4-5-20251001"),
2080 ModelFamily::Haiku
2081 );
2082 }
2083
2084 #[test]
2085 fn model_family_normalizes_bedrock_identifiers() {
2086 let registry = embedded_only();
2087 assert_eq!(
2088 registry.get_model_family("us.anthropic.claude-3-7-sonnet-20250219-v1:0"),
2089 ModelFamily::Sonnet
2090 );
2091 }
2092
2093 #[test]
2094 fn model_family_falls_back_to_the_raw_identifier_when_unregistered() {
2095 let registry = embedded_only();
2096 assert!(registry.get_model_spec("claude-opus-6").is_none());
2100 assert_eq!(
2101 registry.get_model_family("claude-opus-6"),
2102 ModelFamily::Opus
2103 );
2104 }
2105
2106 #[test]
2107 fn model_family_is_unknown_for_unrecognized_identifiers() {
2108 let registry = embedded_only();
2109 assert_eq!(registry.get_model_family("gpt-5.2"), ModelFamily::Unknown);
2110 assert_eq!(
2111 registry.get_model_family("some-future-model"),
2112 ModelFamily::Unknown
2113 );
2114 }
2115
2116 #[test]
2117 fn model_family_glyph_and_label_cover_every_variant() {
2118 assert_eq!(ModelFamily::Fable.glyph(), "π ");
2119 assert_eq!(ModelFamily::Opus.glyph(), "π‘");
2120 assert_eq!(ModelFamily::Sonnet.glyph(), "π’");
2121 assert_eq!(ModelFamily::Haiku.glyph(), "π΅");
2122 assert_eq!(ModelFamily::Unknown.glyph(), "βͺ");
2123 assert_eq!(ModelFamily::Fable.label(), "Fable");
2124 assert_eq!(ModelFamily::Opus.label(), "Opus");
2125 assert_eq!(ModelFamily::Sonnet.label(), "Sonnet");
2126 assert_eq!(ModelFamily::Haiku.label(), "Haiku");
2127 assert_eq!(ModelFamily::Unknown.label(), "Claude");
2128 }
2129}