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-5",
1308 "claude-sonnet-5",
1309 ] {
1310 let spec = registry
1311 .get_model_spec(id)
1312 .unwrap_or_else(|| panic!("{id} must be registered"));
1313 assert_eq!(spec.input_context, 1_000_000, "{id} context window");
1314 assert_eq!(spec.max_output_tokens, 128_000, "{id} max output");
1315 assert!(
1316 spec.beta_headers.is_empty(),
1317 "{id} exposes 1M natively and must not rely on a context beta header"
1318 );
1319 assert!(spec.input_token_price.is_some(), "{id} input price");
1320 assert!(spec.output_token_price.is_some(), "{id} output price");
1321 assert!(!spec.legacy, "{id} is current, not legacy");
1322 }
1323
1324 let sonnet5 = registry.get_model_spec("claude-sonnet-5").unwrap();
1327 assert_eq!(sonnet5.input_token_price, Some(3.0));
1328 assert_eq!(sonnet5.output_token_price, Some(15.0));
1329 let sonnet55 = registry.get_model_spec("claude-sonnet-5-5").unwrap();
1330 assert_eq!(sonnet55.input_token_price, Some(2.0));
1331 assert_eq!(sonnet55.output_token_price, Some(10.0));
1332 }
1333
1334 #[test]
1338 fn claude_4_6_limits_are_native_not_beta_gated() {
1339 let registry = embedded_only();
1340
1341 assert_eq!(registry.get_input_context("claude-opus-4-6"), 1_000_000);
1342 assert_eq!(registry.get_input_context("claude-sonnet-4-6"), 1_000_000);
1343 assert_eq!(registry.get_max_output_tokens("claude-sonnet-4-6"), 128_000);
1344 assert_eq!(registry.get_max_output_tokens("claude-opus-4-6"), 128_000);
1345 }
1346
1347 #[test]
1354 fn undated_aliases_resolve_to_real_limits() {
1355 let registry = embedded_only();
1356
1357 for (alias, dated) in [
1358 ("claude-haiku-4-5", "claude-haiku-4-5-20251001"),
1359 ("claude-sonnet-4-5", "claude-sonnet-4-5-20250929"),
1360 ("claude-opus-4-5", "claude-opus-4-5-20251101"),
1361 ] {
1362 assert_eq!(
1363 registry.get_max_output_tokens(alias),
1364 registry.get_max_output_tokens(dated),
1365 "{alias} must resolve to the same output cap as {dated}, \
1366 not the provider fallback"
1367 );
1368 assert_eq!(
1369 registry.get_input_context(alias),
1370 registry.get_input_context(dated),
1371 "{alias} context window must match {dated}"
1372 );
1373 }
1374
1375 assert_eq!(registry.get_max_output_tokens("claude-haiku-4-5"), 64_000);
1377 }
1378
1379 #[test]
1383 fn retires_records_announced_dates_only() {
1384 let registry = embedded_only();
1385
1386 let opus_41 = registry.get_model_spec("claude-opus-4-1-20250805").unwrap();
1387 assert_eq!(opus_41.retires.as_deref(), Some("2026-08-05"));
1388 assert!(opus_41.legacy);
1389
1390 assert_eq!(
1393 registry
1394 .get_model_spec("claude-3-opus-20240229")
1395 .unwrap()
1396 .retires
1397 .as_deref(),
1398 Some("2026-01-05")
1399 );
1400
1401 for id in ["claude-opus-4-20250514", "claude-sonnet-4-20250514"] {
1403 let spec = registry.get_model_spec(id).unwrap();
1404 assert!(spec.legacy, "{id} is deprecated");
1405 assert!(
1406 spec.retires.is_none(),
1407 "{id} has no announced retirement date; absence must not be \
1408 filled in with a guess"
1409 );
1410 }
1411
1412 assert!(registry
1414 .get_model_spec("claude-sonnet-5")
1415 .unwrap()
1416 .retires
1417 .is_none());
1418 }
1419
1420 #[test]
1421 fn extract_core_model_identifier() {
1422 let registry = embedded_only();
1423
1424 assert_eq!(
1426 registry.extract_core_model_identifier("us.anthropic.claude-3-7-sonnet-20250219-v1:0"),
1427 "claude-3-7-sonnet-20250219"
1428 );
1429
1430 assert_eq!(
1431 registry.extract_core_model_identifier("anthropic.claude-3-haiku-20240307-v1:0"),
1432 "claude-3-haiku-20240307"
1433 );
1434
1435 assert_eq!(
1436 registry.extract_core_model_identifier("claude-3-opus-20240229"),
1437 "claude-3-opus-20240229"
1438 );
1439
1440 assert_eq!(
1441 registry.extract_core_model_identifier("eu.anthropic.claude-sonnet-4-20250514-v2:1"),
1442 "claude-sonnet-4-20250514"
1443 );
1444 }
1445
1446 #[test]
1447 fn beta_header_lookups() {
1448 let registry = embedded_only();
1449
1450 assert_eq!(registry.get_max_output_tokens("claude-opus-4-6"), 128_000);
1453 assert_eq!(registry.get_input_context("claude-opus-4-6"), 1_000_000);
1454
1455 assert_eq!(
1460 registry.get_input_context_with_beta("claude-opus-4-6", "context-1m-2025-08-07"),
1461 1_000_000
1462 );
1463 assert_eq!(
1465 registry.get_max_output_tokens_with_beta("claude-opus-4-6", "context-1m-2025-08-07"),
1466 128_000
1467 );
1468
1469 assert_eq!(
1471 registry.get_max_output_tokens_with_beta(
1472 "claude-3-7-sonnet-20250219",
1473 "output-128k-2025-02-19"
1474 ),
1475 128_000
1476 );
1477
1478 assert_eq!(
1480 registry.get_max_output_tokens("claude-3-7-sonnet-20250219"),
1481 64000
1482 );
1483
1484 let headers = registry.get_beta_headers("claude-opus-4-6");
1486 assert_eq!(headers.len(), 1);
1487 assert_eq!(headers[0].key, "anthropic-beta");
1488 assert_eq!(headers[0].value, "context-1m-2025-08-07");
1489
1490 let headers = registry.get_beta_headers("claude-3-7-sonnet-20250219");
1492 assert_eq!(headers.len(), 2);
1493
1494 let headers = registry.get_beta_headers("claude-3-haiku-20240307");
1496 assert!(headers.is_empty());
1497
1498 let headers = registry.get_beta_headers("unknown-model");
1500 assert!(headers.is_empty());
1501 }
1502
1503 #[test]
1504 fn beta_lookups_for_unknown_model_fall_through_to_provider_defaults() {
1505 let registry = embedded_only();
1506
1507 assert_eq!(
1511 registry
1512 .get_max_output_tokens_with_beta("claude-unknown-model", "context-1m-2025-08-07"),
1513 4096
1514 );
1515 assert_eq!(
1516 registry.get_input_context_with_beta("claude-unknown-model", "context-1m-2025-08-07"),
1517 200_000
1518 );
1519 }
1520
1521 #[test]
1522 fn embedded_models_default_to_embedded_source() {
1523 let registry = embedded_only();
1524 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1525 assert_eq!(spec.source, ModelSource::Embedded);
1526
1527 let provider = registry.get_provider_config("claude").unwrap();
1528 assert_eq!(provider.source, ModelSource::Embedded);
1529 }
1530
1531 #[test]
1532 fn missing_user_and_project_files_fall_through_silently() {
1533 let dir = tempfile::tempdir().unwrap();
1534 let project_path = dir.path().join("missing-project.yaml");
1535 let user_path = dir.path().join("missing-user.yaml");
1536 let registry =
1537 ModelRegistry::load_layered_from_paths(Some(&project_path), Some(&user_path), None)
1538 .unwrap();
1539
1540 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1542 assert_eq!(spec.source, ModelSource::Embedded);
1543 assert_eq!(spec.max_output_tokens, 128_000);
1544 }
1545
1546 #[test]
1547 fn user_layer_overrides_embedded_entry() {
1548 let dir = tempfile::tempdir().unwrap();
1549 let user = write_yaml(
1550 dir.path(),
1551 "user.yaml",
1552 r#"
1553version: "1"
1554models:
1555 - provider: "claude"
1556 model: "Claude Opus 4.6 (custom)"
1557 api_identifier: "claude-opus-4-6"
1558 max_output_tokens: 999999
1559 input_context: 200000
1560 generation: 4.6
1561 tier: "flagship"
1562"#,
1563 );
1564
1565 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1566 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1567 assert_eq!(spec.max_output_tokens, 999_999);
1568 assert_eq!(spec.model, "Claude Opus 4.6 (custom)");
1569 assert_eq!(spec.source, ModelSource::User);
1570 }
1571
1572 #[test]
1573 fn project_layer_takes_precedence_over_user_layer() {
1574 let dir = tempfile::tempdir().unwrap();
1575 let user = write_yaml(
1576 dir.path(),
1577 "user.yaml",
1578 r#"
1579version: "1"
1580models:
1581 - provider: "claude"
1582 model: "From User"
1583 api_identifier: "claude-opus-4-6"
1584 max_output_tokens: 1
1585 input_context: 1
1586 generation: 4.6
1587 tier: "flagship"
1588"#,
1589 );
1590 let project = write_yaml(
1591 dir.path(),
1592 "project.yaml",
1593 r#"
1594version: "1"
1595models:
1596 - provider: "claude"
1597 model: "From Project"
1598 api_identifier: "claude-opus-4-6"
1599 max_output_tokens: 2
1600 input_context: 2
1601 generation: 4.6
1602 tier: "flagship"
1603"#,
1604 );
1605
1606 let registry =
1607 ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
1608 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1609 assert_eq!(spec.model, "From Project");
1610 assert_eq!(spec.max_output_tokens, 2);
1611 assert_eq!(spec.source, ModelSource::Project);
1612 }
1613
1614 #[test]
1615 fn additive_user_entry_is_appended() {
1616 let dir = tempfile::tempdir().unwrap();
1617 let user = write_yaml(
1618 dir.path(),
1619 "user.yaml",
1620 r#"
1621version: "1"
1622models:
1623 - provider: "claude"
1624 model: "Claude Custom Future"
1625 api_identifier: "claude-future-9000"
1626 max_output_tokens: 250000
1627 input_context: 5000000
1628 generation: 9.0
1629 tier: "flagship"
1630"#,
1631 );
1632
1633 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1634 let spec = registry.get_model_spec("claude-future-9000").unwrap();
1635 assert_eq!(spec.max_output_tokens, 250_000);
1636 assert_eq!(spec.input_context, 5_000_000);
1637 assert_eq!(spec.source, ModelSource::User);
1638
1639 let opus = registry.get_model_spec("claude-opus-4-6").unwrap();
1641 assert_eq!(opus.source, ModelSource::Embedded);
1642 }
1643
1644 #[test]
1645 fn provider_fields_can_be_partially_overridden() {
1646 let dir = tempfile::tempdir().unwrap();
1647 let user = write_yaml(
1650 dir.path(),
1651 "user.yaml",
1652 r#"
1653version: "1"
1654providers:
1655 claude:
1656 default_model: "claude-opus-4-6"
1657"#,
1658 );
1659
1660 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1661 let claude = registry.get_provider_config("claude").unwrap();
1662 assert_eq!(claude.default_model, "claude-opus-4-6");
1663 assert_eq!(claude.name, "Anthropic Claude");
1665 assert_eq!(claude.api_base, "https://api.anthropic.com/v1");
1666 assert!(claude.tiers.contains_key("flagship"));
1667 assert_eq!(claude.source, ModelSource::User);
1669 }
1670
1671 #[test]
1672 fn malformed_user_yaml_logs_and_falls_through() {
1673 let dir = tempfile::tempdir().unwrap();
1674 let user = write_yaml(
1675 dir.path(),
1676 "user.yaml",
1677 "this: is: definitely: not: valid: yaml: [unbalanced",
1678 );
1679
1680 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1681 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1683 assert_eq!(spec.source, ModelSource::Embedded);
1684 assert_eq!(spec.max_output_tokens, 128_000);
1685 }
1686
1687 #[test]
1688 fn override_path_short_circuits_user_and_project() {
1689 let dir = tempfile::tempdir().unwrap();
1690 let user = write_yaml(
1691 dir.path(),
1692 "user.yaml",
1693 r#"
1694version: "1"
1695models:
1696 - provider: "claude"
1697 model: "From User"
1698 api_identifier: "claude-opus-4-6"
1699 max_output_tokens: 1
1700 input_context: 1
1701 generation: 4.6
1702 tier: "flagship"
1703"#,
1704 );
1705 let project = write_yaml(
1706 dir.path(),
1707 "project.yaml",
1708 r#"
1709version: "1"
1710models:
1711 - provider: "claude"
1712 model: "From Project"
1713 api_identifier: "claude-opus-4-6"
1714 max_output_tokens: 2
1715 input_context: 2
1716 generation: 4.6
1717 tier: "flagship"
1718"#,
1719 );
1720 let override_file = write_yaml(
1721 dir.path(),
1722 "override.yaml",
1723 r#"
1724version: "1"
1725models:
1726 - provider: "claude"
1727 model: "From Override"
1728 api_identifier: "claude-opus-4-6"
1729 max_output_tokens: 3
1730 input_context: 3
1731 generation: 4.6
1732 tier: "flagship"
1733"#,
1734 );
1735
1736 let registry = ModelRegistry::load_layered_from_paths(
1737 Some(&project),
1738 Some(&user),
1739 Some(&override_file),
1740 )
1741 .unwrap();
1742 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1743 assert_eq!(spec.model, "From Override");
1744 assert_eq!(spec.max_output_tokens, 3);
1745 assert_eq!(spec.source, ModelSource::Override);
1746 }
1747
1748 #[test]
1749 fn missing_override_path_falls_back_to_embedded() {
1750 let dir = tempfile::tempdir().unwrap();
1751 let missing = dir.path().join("does-not-exist.yaml");
1752 let registry = ModelRegistry::load_layered_from_paths(None, None, Some(&missing)).unwrap();
1753 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1754 assert_eq!(spec.source, ModelSource::Embedded);
1755 }
1756
1757 #[test]
1758 fn version_mismatch_is_warned_not_fatal() {
1759 let dir = tempfile::tempdir().unwrap();
1760 let user = write_yaml(
1761 dir.path(),
1762 "user.yaml",
1763 r#"
1764version: "9999"
1765models:
1766 - provider: "claude"
1767 model: "From Future"
1768 api_identifier: "claude-future-9000"
1769 max_output_tokens: 1
1770 input_context: 1
1771 generation: 9.0
1772 tier: "flagship"
1773"#,
1774 );
1775 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1776 assert!(registry.get_model_spec("claude-future-9000").is_some());
1778 }
1779
1780 #[test]
1781 fn missing_version_is_accepted() {
1782 let dir = tempfile::tempdir().unwrap();
1783 let user = write_yaml(
1784 dir.path(),
1785 "user.yaml",
1786 r#"
1787models:
1788 - provider: "claude"
1789 model: "Versionless"
1790 api_identifier: "claude-versionless"
1791 max_output_tokens: 1
1792 input_context: 1
1793 generation: 1.0
1794 tier: "flagship"
1795"#,
1796 );
1797 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1798 assert!(registry.get_model_spec("claude-versionless").is_some());
1799 }
1800
1801 #[test]
1802 fn model_entry_without_api_identifier_is_skipped() {
1803 let dir = tempfile::tempdir().unwrap();
1804 let user = write_yaml(
1805 dir.path(),
1806 "user.yaml",
1807 r#"
1808version: "1"
1809models:
1810 - provider: "claude"
1811 model: "No Id"
1812 max_output_tokens: 1
1813 input_context: 1
1814 generation: 1.0
1815 tier: "flagship"
1816"#,
1817 );
1818 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1819 let opus = registry.get_model_spec("claude-opus-4-6").unwrap();
1821 assert_eq!(opus.source, ModelSource::Embedded);
1822 }
1823
1824 #[test]
1825 fn model_source_display() {
1826 assert_eq!(ModelSource::Embedded.to_string(), "embedded");
1827 assert_eq!(ModelSource::User.to_string(), "user");
1828 assert_eq!(ModelSource::Project.to_string(), "project");
1829 assert_eq!(ModelSource::Override.to_string(), "override");
1830 }
1831
1832 #[test]
1833 fn embedded_yaml_must_not_be_malformed() {
1834 let layers = [(ModelSource::Embedded, "::: not yaml :::".to_string())];
1836 let result = ModelRegistry::from_layers(&layers);
1837 assert!(result.is_err());
1838 }
1839
1840 #[test]
1841 fn user_layer_with_scalar_top_level_returns_error() {
1842 let dir = tempfile::tempdir().unwrap();
1846 let user = write_yaml(dir.path(), "user.yaml", "\"just a string\"\n");
1847 let result = ModelRegistry::load_layered_from_paths(None, Some(&user), None);
1848 assert!(result.is_err());
1849 }
1850
1851 #[test]
1852 fn user_layer_with_non_sequence_models_returns_error() {
1853 let dir = tempfile::tempdir().unwrap();
1857 let user = write_yaml(
1858 dir.path(),
1859 "user.yaml",
1860 r#"
1861version: "1"
1862models: 42
1863"#,
1864 );
1865 let result = ModelRegistry::load_layered_from_paths(None, Some(&user), None);
1866 assert!(result.is_err());
1867 }
1868
1869 #[test]
1870 fn user_layer_with_non_mapping_providers_returns_error() {
1871 let dir = tempfile::tempdir().unwrap();
1874 let user = write_yaml(
1875 dir.path(),
1876 "user.yaml",
1877 r#"
1878version: "1"
1879providers: 42
1880"#,
1881 );
1882 let result = ModelRegistry::load_layered_from_paths(None, Some(&user), None);
1883 assert!(result.is_err());
1884 }
1885
1886 #[test]
1887 fn deep_merge_inserts_new_keys_into_existing_mapping() {
1888 let dir = tempfile::tempdir().unwrap();
1892 let user = write_yaml(
1893 dir.path(),
1894 "user.yaml",
1895 r#"
1896version: "1"
1897providers:
1898 claude:
1899 tiers:
1900 experimental:
1901 description: "Experimental tier"
1902 use_cases: ["bleeding edge"]
1903"#,
1904 );
1905 let registry = ModelRegistry::load_layered_from_paths(None, Some(&user), None).unwrap();
1906 let claude = registry.get_provider_config("claude").unwrap();
1907 assert!(claude.tiers.contains_key("flagship"));
1909 assert!(claude.tiers.contains_key("balanced"));
1910 assert!(claude.tiers.contains_key("fast"));
1911 let experimental = claude.tiers.get("experimental").unwrap();
1913 assert_eq!(experimental.description, "Experimental tier");
1914 assert_eq!(experimental.use_cases, vec!["bleeding edge".to_string()]);
1915 }
1916
1917 #[test]
1918 #[cfg(unix)]
1919 fn user_path_pointing_at_a_directory_logs_and_falls_through() {
1920 let dir = tempfile::tempdir().unwrap();
1923 let bogus = dir.path().join("models.yaml");
1924 std::fs::create_dir(&bogus).unwrap();
1925 let registry = ModelRegistry::load_layered_from_paths(None, Some(&bogus), None).unwrap();
1926 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1927 assert_eq!(spec.source, ModelSource::Embedded);
1928 }
1929
1930 #[test]
1931 #[cfg(unix)]
1932 fn override_path_pointing_at_a_directory_warns_and_falls_through() {
1933 let dir = tempfile::tempdir().unwrap();
1934 let bogus = dir.path().join("override.yaml");
1935 std::fs::create_dir(&bogus).unwrap();
1936 let registry = ModelRegistry::load_layered_from_paths(None, None, Some(&bogus)).unwrap();
1937 let spec = registry.get_model_spec("claude-opus-4-6").unwrap();
1938 assert_eq!(spec.source, ModelSource::Embedded);
1939 }
1940
1941 #[test]
1942 fn project_layer_recovers_after_user_replaces_top_level_with_scalar() {
1943 let dir = tempfile::tempdir().unwrap();
1949 let user = write_yaml(dir.path(), "user.yaml", "\"junk\"\n");
1950 let project = write_yaml(
1951 dir.path(),
1952 "project.yaml",
1953 r#"
1954version: "1"
1955models:
1956 - provider: "claude"
1957 model: "Project Rescue"
1958 api_identifier: "claude-rescue"
1959 max_output_tokens: 1
1960 input_context: 1
1961 generation: 1.0
1962 tier: "flagship"
1963providers:
1964 custom-provider:
1965 name: "Custom"
1966 api_base: "https://example.invalid"
1967 default_model: "custom-default"
1968 tiers: {}
1969 defaults:
1970 max_output_tokens: 100
1971 input_context: 1000
1972"#,
1973 );
1974 let registry =
1975 ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
1976 let spec = registry.get_model_spec("claude-rescue").unwrap();
1978 assert_eq!(spec.source, ModelSource::Project);
1979 }
1980
1981 #[test]
1982 fn project_layer_recovers_after_user_replaces_models_with_scalar() {
1983 let dir = tempfile::tempdir().unwrap();
1987 let user = write_yaml(
1988 dir.path(),
1989 "user.yaml",
1990 r#"
1991version: "1"
1992models: 42
1993"#,
1994 );
1995 let project = write_yaml(
1996 dir.path(),
1997 "project.yaml",
1998 r#"
1999version: "1"
2000models:
2001 - provider: "claude"
2002 model: "Project Rescue"
2003 api_identifier: "claude-rescue"
2004 max_output_tokens: 1
2005 input_context: 1
2006 generation: 1.0
2007 tier: "flagship"
2008"#,
2009 );
2010 let registry =
2011 ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
2012 let spec = registry.get_model_spec("claude-rescue").unwrap();
2013 assert_eq!(spec.source, ModelSource::Project);
2014 }
2015
2016 #[test]
2017 fn project_layer_recovers_after_user_replaces_providers_with_scalar() {
2018 let dir = tempfile::tempdir().unwrap();
2021 let user = write_yaml(
2022 dir.path(),
2023 "user.yaml",
2024 r#"
2025version: "1"
2026providers: 42
2027"#,
2028 );
2029 let project = write_yaml(
2030 dir.path(),
2031 "project.yaml",
2032 r#"
2033version: "1"
2034providers:
2035 custom-provider:
2036 name: "Custom"
2037 api_base: "https://example.invalid"
2038 default_model: "custom-default"
2039 tiers: {}
2040 defaults:
2041 max_output_tokens: 100
2042 input_context: 1000
2043"#,
2044 );
2045 let registry =
2046 ModelRegistry::load_layered_from_paths(Some(&project), Some(&user), None).unwrap();
2047 let provider = registry.get_provider_config("custom-provider").unwrap();
2048 assert_eq!(provider.name, "Custom");
2049 assert_eq!(provider.source, ModelSource::Project);
2050 }
2051
2052 #[test]
2053 fn empty_omni_dev_models_yaml_env_var_is_ignored() {
2054 let resolved: Option<PathBuf> = Some(String::new())
2058 .filter(|s| !s.is_empty())
2059 .map(PathBuf::from);
2060 assert!(resolved.is_none());
2061 let resolved: Option<PathBuf> = Some("/some/path".to_string())
2062 .filter(|s| !s.is_empty())
2063 .map(PathBuf::from);
2064 assert_eq!(resolved.as_deref(), Some(Path::new("/some/path")));
2065 }
2066
2067 #[test]
2068 fn model_family_classifies_registered_claude_identifiers() {
2069 let registry = embedded_only();
2070 assert_eq!(
2071 registry.get_model_family("claude-fable-5"),
2072 ModelFamily::Fable
2073 );
2074 assert_eq!(
2075 registry.get_model_family("claude-opus-4-8"),
2076 ModelFamily::Opus
2077 );
2078 assert_eq!(
2079 registry.get_model_family("claude-sonnet-5"),
2080 ModelFamily::Sonnet
2081 );
2082 assert_eq!(
2083 registry.get_model_family("claude-haiku-4-5-20251001"),
2084 ModelFamily::Haiku
2085 );
2086 }
2087
2088 #[test]
2089 fn model_family_normalizes_bedrock_identifiers() {
2090 let registry = embedded_only();
2091 assert_eq!(
2092 registry.get_model_family("us.anthropic.claude-3-7-sonnet-20250219-v1:0"),
2093 ModelFamily::Sonnet
2094 );
2095 }
2096
2097 #[test]
2098 fn model_family_falls_back_to_the_raw_identifier_when_unregistered() {
2099 let registry = embedded_only();
2100 assert!(registry.get_model_spec("claude-opus-6").is_none());
2104 assert_eq!(
2105 registry.get_model_family("claude-opus-6"),
2106 ModelFamily::Opus
2107 );
2108 }
2109
2110 #[test]
2111 fn model_family_is_unknown_for_unrecognized_identifiers() {
2112 let registry = embedded_only();
2113 assert_eq!(registry.get_model_family("gpt-5.2"), ModelFamily::Unknown);
2114 assert_eq!(
2115 registry.get_model_family("some-future-model"),
2116 ModelFamily::Unknown
2117 );
2118 }
2119
2120 #[test]
2121 fn model_family_glyph_and_label_cover_every_variant() {
2122 assert_eq!(ModelFamily::Fable.glyph(), "π ");
2123 assert_eq!(ModelFamily::Opus.glyph(), "π‘");
2124 assert_eq!(ModelFamily::Sonnet.glyph(), "π’");
2125 assert_eq!(ModelFamily::Haiku.glyph(), "π΅");
2126 assert_eq!(ModelFamily::Unknown.glyph(), "βͺ");
2127 assert_eq!(ModelFamily::Fable.label(), "Fable");
2128 assert_eq!(ModelFamily::Opus.label(), "Opus");
2129 assert_eq!(ModelFamily::Sonnet.label(), "Sonnet");
2130 assert_eq!(ModelFamily::Haiku.label(), "Haiku");
2131 assert_eq!(ModelFamily::Unknown.label(), "Claude");
2132 }
2133}