1use crate::route::{RoutePolicy, RouteTier};
16use crate::secret::SecretRef;
17use serde::{Deserialize, Serialize};
18use std::collections::BTreeMap;
19use std::path::{Path, PathBuf};
20
21#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
22pub struct ModelEntry {
23 #[serde(default)]
24 pub provider: String,
25 #[serde(default, skip_serializing_if = "Option::is_none")]
37 pub vendor: Option<String>,
38 #[serde(default)]
39 pub model: String,
40 #[serde(default, skip_serializing_if = "Option::is_none")]
41 pub base_url: Option<String>,
42 #[serde(default, skip_serializing_if = "Option::is_none")]
43 pub secret: Option<SecretRef>,
44 #[serde(default, skip_serializing_if = "Vec::is_empty")]
45 pub capabilities: Vec<String>,
46 #[serde(default, skip_serializing_if = "serde_json::Value::is_null")]
47 pub params: serde_json::Value,
48 #[serde(default, skip_serializing_if = "Option::is_none")]
51 pub tier: Option<RouteTier>,
52 #[serde(default, skip_serializing_if = "Option::is_none")]
55 pub cost_per_1k_tokens: Option<f64>,
56 #[serde(default, skip_serializing_if = "Option::is_none")]
59 pub input_cost_per_1k: Option<f64>,
60 #[serde(default, skip_serializing_if = "Option::is_none")]
63 pub output_cost_per_1k: Option<f64>,
64 #[serde(default, skip_serializing_if = "Option::is_none")]
66 pub context_window: Option<u64>,
67 #[serde(default, skip_serializing_if = "Option::is_none")]
74 pub priced_at: Option<chrono::DateTime<chrono::Utc>>,
75}
76
77fn vendor_label_of_url(base_url: Option<&str>) -> Option<String> {
82 let host = base_url?
83 .split("//")
84 .nth(1)
85 .unwrap_or(base_url?)
86 .split(['/', ':'])
87 .next()
88 .unwrap_or("");
89 let label = host.strip_prefix("api.").unwrap_or(host);
90 let first = label.split('.').next().unwrap_or("");
91 (!first.is_empty()).then(|| first.to_string())
92}
93
94impl ModelEntry {
95 pub fn effective_costs(&self) -> (Option<f64>, Option<f64>) {
100 let output = self.output_cost_per_1k.or(self.cost_per_1k_tokens);
101 let input = self.input_cost_per_1k.or(self.cost_per_1k_tokens);
102 (input, output)
103 }
104
105 pub fn vendor_candidates(&self) -> Vec<String> {
116 let mut out: Vec<String> = Vec::with_capacity(3);
117 let mut push = |v: &str| {
118 if !v.is_empty() && !out.iter().any(|e| e == v) {
119 out.push(v.to_string());
120 }
121 };
122 if let Some(v) = self.vendor.as_deref() {
123 push(v);
124 }
125 if let Some(label) = vendor_label_of_url(self.base_url.as_deref()) {
126 push(&label);
127 }
128 push(&self.provider);
129 out
130 }
131
132 pub fn is_priced(&self) -> bool {
134 let (input, output) = self.effective_costs();
135 input.is_some() || output.is_some()
136 }
137
138 pub fn stamp_priced_at(&mut self, now: chrono::DateTime<chrono::Utc>) {
142 if self.is_priced() && self.priced_at.is_none_or(|prev| prev < now) {
143 self.priced_at = Some(now);
144 }
145 }
146
147 pub fn price_age(&self, now: chrono::DateTime<chrono::Utc>) -> Option<chrono::TimeDelta> {
149 self.priced_at.map(|at| now - at)
150 }
151}
152
153#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
154pub struct RoleEntry {
155 pub primary: String,
157 #[serde(default, skip_serializing_if = "Option::is_none")]
159 pub fallback: Option<String>,
160 #[serde(default, skip_serializing_if = "Option::is_none")]
162 pub cost_budget_per_day_usd: Option<f64>,
163 #[serde(default)]
165 pub privacy_local_only: bool,
166 #[serde(default, skip_serializing_if = "Option::is_none")]
169 pub route_policy: Option<RoutePolicy>,
170}
171
172#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
173pub struct ModelRegistry {
174 pub schema_version: u32,
175 #[serde(default)]
176 pub models: BTreeMap<String, ModelEntry>,
177 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
178 pub roles: BTreeMap<String, RoleEntry>,
179}
180
181impl Default for ModelRegistry {
182 fn default() -> Self {
183 Self {
184 schema_version: 1,
185 models: BTreeMap::new(),
186 roles: BTreeMap::new(),
187 }
188 }
189}
190
191impl ModelRegistry {
192 pub fn load_from(path: &Path) -> anyhow::Result<Self> {
193 if !path.exists() {
194 return Ok(Self::default());
195 }
196 let body = std::fs::read_to_string(path)?;
197 if body.trim().is_empty() {
198 return Ok(Self::default());
199 }
200 Ok(serde_yaml_ng::from_str(&body)?)
201 }
202
203 pub fn save_to(&self, path: &Path) -> anyhow::Result<()> {
204 if let Some(parent) = path.parent() {
205 std::fs::create_dir_all(parent)?;
206 }
207 let body = serde_yaml_ng::to_string(self)?;
208 let tmp = path.with_extension("yaml.tmp");
209 std::fs::write(&tmp, body)?;
210 std::fs::rename(&tmp, path)?;
211 Ok(())
212 }
213
214 pub fn default_path() -> anyhow::Result<PathBuf> {
215 if let Ok(p) = std::env::var("MUR_HOME")
218 && !p.is_empty()
219 {
220 return Ok(PathBuf::from(p).join("models.yaml"));
221 }
222 let home = dirs::home_dir().ok_or_else(|| anyhow::anyhow!("no home dir"))?;
223 Ok(home.join(".mur/models.yaml"))
224 }
225
226 pub fn resolve_role(&self, role: &str) -> Option<&str> {
229 let entry = self.roles.get(role)?;
230 if self.models.contains_key(&entry.primary) {
231 return Some(&entry.primary);
232 }
233 if let Some(fb) = &entry.fallback
235 && self.models.contains_key(fb)
236 {
237 return Some(fb);
238 }
239 None
241 }
242}
243
244use crate::agent::AgentProfile;
245use crate::config::{DEFAULT_ROUTING_THRESHOLD, ModelSwitchConfig, RoutingConfig};
246
247pub fn resolve_model_refs(
252 profile: &AgentProfile,
253 cfg: &ModelSwitchConfig,
254 routed_primary: Option<String>,
255) -> Vec<String> {
256 let primary = routed_primary
257 .or_else(|| profile.model_ref.clone())
258 .or_else(|| cfg.default.clone());
259 let chain = if !profile.fallback_chain.is_empty() {
260 profile.fallback_chain.clone()
261 } else {
262 cfg.fallback_chain.clone()
263 };
264 let mut out: Vec<String> = Vec::new();
265 if let Some(p) = primary {
266 out.push(p);
267 }
268 for r in chain {
269 if !out.contains(&r) {
270 out.push(r);
271 }
272 }
273 out
274}
275
276pub fn choose_by_difficulty(est_input_tokens: u32, r: &RoutingConfig) -> Option<String> {
280 let threshold = r
281 .threshold_input_tokens
282 .unwrap_or(DEFAULT_ROUTING_THRESHOLD);
283 match (r.cheap.as_ref(), r.frontier.as_ref()) {
284 (Some(cheap), Some(frontier)) => Some(if est_input_tokens > threshold {
285 frontier.clone()
286 } else {
287 cheap.clone()
288 }),
289 _ => None,
290 }
291}
292
293pub const CAP_CHAT: &str = "chat";
297pub const CAP_TOOLS: &str = "tools";
298pub const CAP_VISION: &str = "vision";
299
300#[derive(Debug, Clone, Copy, PartialEq, Eq)]
304pub enum Requirement {
305 Vision,
307 Tools,
309}
310
311impl Requirement {
312 pub fn capability(self) -> &'static str {
314 match self {
315 Requirement::Vision => CAP_VISION,
316 Requirement::Tools => CAP_TOOLS,
317 }
318 }
319
320 fn permitted_when_undeclared(self) -> bool {
340 match self {
341 Requirement::Vision => false,
342 Requirement::Tools => true,
343 }
344 }
345}
346
347pub fn satisfies(e: &ModelEntry, reqs: &[Requirement]) -> bool {
355 let chat_capable = e.capabilities.is_empty() || e.capabilities.iter().any(|c| c == CAP_CHAT);
356 if !chat_capable {
357 return false;
358 }
359 reqs.iter().all(|r| {
360 if e.capabilities.is_empty() {
361 r.permitted_when_undeclared()
362 } else {
363 e.capabilities.iter().any(|c| c == r.capability())
364 }
365 })
366}
367
368pub fn pick_cheap_model(
372 reg: &ModelRegistry,
373 exclude: Option<&str>,
374 reqs: &[Requirement],
375) -> Option<String> {
376 reg.models
377 .iter()
378 .filter(|(k, _)| exclude != Some(k.as_str()))
379 .filter(|(_, e)| satisfies(e, reqs))
380 .filter_map(|(k, e)| {
381 let (input, output) = e.effective_costs();
385 output.or(input).map(|c| (c, k.clone()))
386 })
387 .min_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal))
388 .map(|(_, k)| k)
389}
390
391#[cfg(test)]
392mod tests {
393
394 #[test]
395 fn vendor_candidates_prefer_the_recorded_vendor_then_the_host_then_provider() {
396 let e = ModelEntry {
398 provider: "openai".into(),
399 vendor: Some("deepseek".into()),
400 base_url: Some("https://api.deepseek.com/v1".into()),
401 ..Default::default()
402 };
403 assert_eq!(e.vendor_candidates(), vec!["deepseek", "openai"]);
404
405 let legacy = ModelEntry {
408 provider: "openai".into(),
409 base_url: Some("https://api.deepseek.com/v1".into()),
410 ..Default::default()
411 };
412 assert_eq!(legacy.vendor_candidates(), vec!["deepseek", "openai"]);
413
414 let bare = ModelEntry {
416 provider: "anthropic".into(),
417 ..Default::default()
418 };
419 assert_eq!(bare.vendor_candidates(), vec!["anthropic"]);
420
421 let same = ModelEntry {
423 provider: "openai".into(),
424 base_url: Some("https://api.openai.com/v1".into()),
425 ..Default::default()
426 };
427 assert_eq!(same.vendor_candidates(), vec!["openai"]);
428 }
429
430 #[test]
431 fn vendor_is_omitted_from_yaml_when_absent_and_round_trips_when_set() {
432 let bare = ModelEntry {
433 provider: "anthropic".into(),
434 model: "claude-opus-5".into(),
435 ..Default::default()
436 };
437 let y = serde_yaml_ng::to_string(&bare).unwrap();
438 assert!(!y.contains("vendor"), "{y}");
439
440 let tagged = ModelEntry {
441 provider: "openai".into(),
442 vendor: Some("groq".into()),
443 model: "llama-3.3".into(),
444 ..Default::default()
445 };
446 let y = serde_yaml_ng::to_string(&tagged).unwrap();
447 let back: ModelEntry = serde_yaml_ng::from_str(&y).unwrap();
448 assert_eq!(back.vendor.as_deref(), Some("groq"));
449 }
450 use super::*;
451
452 #[test]
453 fn parses_full_registry() {
454 let yaml = r#"
455schema_version: 1
456models:
457 anthropic_opus_4_7:
458 provider: anthropic
459 model: claude-opus-4-7
460 secret: env:ANTHROPIC_API_KEY
461 capabilities: [chat, tools]
462 ollama_llama3:
463 provider: ollama
464 model: llama3.2:3b
465 base_url: http://127.0.0.1:11434
466"#;
467 let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
468 assert_eq!(r.schema_version, 1);
469 assert_eq!(r.models.len(), 2);
470 let opus = r.models.get("anthropic_opus_4_7").unwrap();
471 assert_eq!(opus.provider, "anthropic");
472 assert_eq!(
473 opus.secret,
474 Some(SecretRef::Env("ANTHROPIC_API_KEY".into()))
475 );
476 assert!(r.models["ollama_llama3"].secret.is_none());
477 }
478
479 #[test]
480 fn round_trip_preserves_shape() {
481 let mut r = ModelRegistry::default();
482 r.models.insert(
483 "foo".into(),
484 ModelEntry {
485 provider: "anthropic".into(),
486 model: "claude-opus-4-7".into(),
487 base_url: None,
488 secret: Some(SecretRef::Keychain {
489 service: "mur".into(),
490 account: "anthropic".into(),
491 }),
492 capabilities: vec!["chat".into()],
493 params: serde_json::Value::Null,
494 tier: None,
495 cost_per_1k_tokens: None,
496 input_cost_per_1k: None,
497 output_cost_per_1k: None,
498 context_window: None,
499 priced_at: None,
500 ..Default::default()
501 },
502 );
503 let s = serde_yaml_ng::to_string(&r).unwrap();
504 let parsed: ModelRegistry = serde_yaml_ng::from_str(&s).unwrap();
505 assert_eq!(r, parsed);
506 }
507
508 #[test]
509 fn rejects_unknown_secret_scheme() {
510 let yaml = r#"
511schema_version: 1
512models:
513 bad:
514 provider: x
515 model: y
516 secret: bogus:value
517"#;
518 let r: Result<ModelRegistry, _> = serde_yaml_ng::from_str(yaml);
519 assert!(r.is_err(), "should reject unknown scheme");
520 }
521
522 #[test]
523 fn test_registry_roundtrip_with_roles() {
524 let yaml = r#"
525schema_version: 1
526models:
527 haiku:
528 provider: anthropic
529 model: claude-haiku-4-5
530roles:
531 reflector:
532 primary: haiku
533 fallback: null
534 cost_budget_per_day_usd: 0.5
535"#;
536 let reg: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
537 assert_eq!(reg.roles["reflector"].primary, "haiku");
538 let back = serde_yaml_ng::to_string(®).unwrap();
539 let reg2: ModelRegistry = serde_yaml_ng::from_str(&back).unwrap();
540 assert_eq!(reg, reg2);
541 }
542
543 #[test]
544 fn test_resolve_role_primary() {
545 let mut reg = ModelRegistry::default();
546 reg.models.insert(
547 "haiku".into(),
548 ModelEntry {
549 provider: "anthropic".into(),
550 model: "claude-haiku-4-5".into(),
551 base_url: None,
552 secret: None,
553 capabilities: vec![],
554 params: serde_json::Value::Null,
555 tier: None,
556 cost_per_1k_tokens: None,
557 input_cost_per_1k: None,
558 output_cost_per_1k: None,
559 context_window: None,
560 priced_at: None,
561 ..Default::default()
562 },
563 );
564 reg.roles.insert(
565 "reflector".into(),
566 RoleEntry {
567 primary: "haiku".into(),
568 fallback: None,
569 ..Default::default()
570 },
571 );
572 assert_eq!(reg.resolve_role("reflector"), Some("haiku"));
573 }
574
575 #[test]
576 fn test_resolve_role_fallback() {
577 let mut reg = ModelRegistry::default();
578 reg.models.insert(
579 "haiku".into(),
580 ModelEntry {
581 provider: "anthropic".into(),
582 model: "claude-haiku-4-5".into(),
583 base_url: None,
584 secret: None,
585 capabilities: vec![],
586 params: serde_json::Value::Null,
587 tier: None,
588 cost_per_1k_tokens: None,
589 input_cost_per_1k: None,
590 output_cost_per_1k: None,
591 context_window: None,
592 priced_at: None,
593 ..Default::default()
594 },
595 );
596 reg.roles.insert(
597 "reflector".into(),
598 RoleEntry {
599 primary: "nonexistent".into(),
600 fallback: Some("haiku".into()),
601 ..Default::default()
602 },
603 );
604 assert_eq!(reg.resolve_role("reflector"), Some("haiku"));
605 }
606
607 #[test]
608 fn test_resolve_role_none() {
609 let reg = ModelRegistry::default();
610 assert_eq!(reg.resolve_role("reflector"), None);
611 }
612
613 #[test]
614 fn model_entry_parses_tier_field() {
615 let yaml = r#"
616schema_version: 1
617models:
618 haiku:
619 provider: anthropic
620 model: claude-haiku-4-5
621 tier: local
622 opus:
623 provider: anthropic
624 model: claude-opus-4-7
625 tier: frontier
626 cost_per_1k_tokens: 0.015
627"#;
628 let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
629 assert_eq!(r.models["haiku"].tier, Some(RouteTier::Local));
630 assert_eq!(r.models["opus"].tier, Some(RouteTier::Frontier));
631 assert_eq!(r.models["opus"].cost_per_1k_tokens, Some(0.015));
632 let mut r2 = ModelRegistry::default();
634 r2.models.insert(
635 "x".into(),
636 ModelEntry {
637 provider: "ollama".into(),
638 model: "llama3".into(),
639 base_url: None,
640 secret: None,
641 capabilities: vec![],
642 params: serde_json::Value::Null,
643 tier: None,
644 cost_per_1k_tokens: None,
645 input_cost_per_1k: None,
646 output_cost_per_1k: None,
647 context_window: None,
648 priced_at: None,
649 ..Default::default()
650 },
651 );
652 let yaml = serde_yaml_ng::to_string(&r2).unwrap();
653 assert!(
654 !yaml.contains("tier:"),
655 "absent tier should not be serialized: {yaml}"
656 );
657 }
658
659 #[test]
660 fn role_entry_parses_route_policy() {
661 let yaml = r#"
662schema_version: 1
663models:
664 haiku:
665 provider: anthropic
666 model: claude-haiku-4-5
667 opus:
668 provider: anthropic
669 model: claude-opus-4-7
670roles:
671 dev:
672 primary: opus
673 route_policy: !force_frontier
674 model_id: opus
675 reflector:
676 primary: haiku
677 route_policy: prefer_local
678 curator:
679 primary: haiku
680 route_policy: force_local
681 chat:
682 primary: haiku
683"#;
684 let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
685 assert_eq!(
686 r.roles["dev"].route_policy,
687 Some(RoutePolicy::ForceFrontier {
688 model_id: "opus".into()
689 })
690 );
691 assert_eq!(
692 r.roles["reflector"].route_policy,
693 Some(RoutePolicy::PreferLocal)
694 );
695 assert_eq!(
696 r.roles["curator"].route_policy,
697 Some(RoutePolicy::ForceLocal)
698 );
699 assert_eq!(r.roles["chat"].route_policy, None);
700 }
701
702 #[test]
703 fn parses_split_cost_fields() {
704 let yaml = r#"
705schema_version: 1
706models:
707 opus:
708 provider: anthropic
709 model: claude-opus-4-8
710 input_cost_per_1k: 0.005
711 output_cost_per_1k: 0.025
712 context_window: 200000
713"#;
714 let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
715 let e = r.models.get("opus").unwrap();
716 assert_eq!(e.input_cost_per_1k, Some(0.005));
717 assert_eq!(e.output_cost_per_1k, Some(0.025));
718 assert_eq!(e.context_window, Some(200_000));
719 }
720
721 #[test]
722 fn default_model_entry_is_empty() {
723 let e = ModelEntry::default();
724 assert!(e.provider.is_empty());
725 assert_eq!(e.input_cost_per_1k, None);
726 assert_eq!(e.output_cost_per_1k, None);
727 assert_eq!(e.context_window, None);
728 }
729
730 #[test]
731 fn effective_costs_fallback_matrix() {
732 let mut e = ModelEntry {
734 cost_per_1k_tokens: Some(0.01),
735 ..Default::default()
736 };
737 assert_eq!(e.effective_costs(), (Some(0.01), Some(0.01)));
738
739 e = ModelEntry {
741 input_cost_per_1k: Some(0.005),
742 output_cost_per_1k: Some(0.025),
743 ..Default::default()
744 };
745 assert_eq!(e.effective_costs(), (Some(0.005), Some(0.025)));
746
747 e = ModelEntry {
749 cost_per_1k_tokens: Some(0.01),
750 input_cost_per_1k: Some(0.005),
751 output_cost_per_1k: Some(0.025),
752 ..Default::default()
753 };
754 assert_eq!(e.effective_costs(), (Some(0.005), Some(0.025)));
755
756 e = ModelEntry::default();
758 assert_eq!(e.effective_costs(), (None, None));
759 }
760}
761
762#[cfg(test)]
763mod io_tests {
764 use super::*;
765 use tempfile::tempdir;
766
767 #[test]
768 fn load_returns_empty_when_file_missing() {
769 let dir = tempdir().unwrap();
770 let r = ModelRegistry::load_from(&dir.path().join("nope.yaml")).unwrap();
771 assert_eq!(r.models.len(), 0);
772 assert_eq!(r.schema_version, 1);
773 }
774
775 #[test]
776 fn save_then_load_round_trips() {
777 let dir = tempdir().unwrap();
778 let p = dir.path().join("models.yaml");
779 let mut r = ModelRegistry::default();
780 r.models.insert(
781 "x".into(),
782 ModelEntry {
783 provider: "ollama".into(),
784 model: "llama3.2:3b".into(),
785 base_url: None,
786 secret: None,
787 capabilities: vec![],
788 params: serde_json::Value::Null,
789 tier: None,
790 cost_per_1k_tokens: None,
791 input_cost_per_1k: None,
792 output_cost_per_1k: None,
793 context_window: None,
794 priced_at: None,
795 ..Default::default()
796 },
797 );
798 r.save_to(&p).unwrap();
799 let r2 = ModelRegistry::load_from(&p).unwrap();
800 assert_eq!(r, r2);
801 }
802
803 #[test]
804 fn save_uses_atomic_rename() {
805 let dir = tempdir().unwrap();
806 let p = dir.path().join("models.yaml");
807 ModelRegistry::default().save_to(&p).unwrap();
808 let temp = dir.path().join("models.yaml.tmp");
809 assert!(!temp.exists(), "atomic temp left behind");
810 }
811}
812
813#[cfg(test)]
814mod switch_tests {
815 use super::*;
816 use crate::agent::AgentProfile;
817 use crate::config::{ModelSwitchConfig, RoutingConfig};
818
819 fn profile(model_ref: Option<&str>, chain: &[&str]) -> AgentProfile {
820 let mut p = AgentProfile::default_for_tests();
821 p.model_ref = model_ref.map(|s| s.to_string());
822 p.fallback_chain = chain.iter().map(|s| s.to_string()).collect();
823 p
824 }
825
826 #[test]
827 fn per_agent_primary_and_chain_win_over_global() {
828 let cfg = ModelSwitchConfig {
829 default: Some("global_default".into()),
830 fallback_chain: vec!["g1".into(), "g2".into()],
831 ..Default::default()
832 };
833 let p = profile(Some("agent_primary"), &["agent_primary", "agent_fb"]);
834 assert_eq!(
836 resolve_model_refs(&p, &cfg, None),
837 vec!["agent_primary", "agent_fb"]
838 );
839 }
840
841 #[test]
842 fn falls_back_to_global_default_and_chain() {
843 let cfg = ModelSwitchConfig {
844 default: Some("global_default".into()),
845 fallback_chain: vec!["g1".into(), "global_default".into()],
846 ..Default::default()
847 };
848 let p = profile(None, &[]); assert_eq!(
851 resolve_model_refs(&p, &cfg, None),
852 vec!["global_default", "g1"]
853 );
854 }
855
856 #[test]
857 fn routed_primary_overrides_model_ref() {
858 let cfg = ModelSwitchConfig {
859 fallback_chain: vec!["g1".into()],
860 ..Default::default()
861 };
862 let p = profile(Some("agent_primary"), &[]);
863 assert_eq!(
864 resolve_model_refs(&p, &cfg, Some("frontier".into())),
865 vec!["frontier", "g1"]
866 );
867 }
868
869 #[test]
870 fn no_config_no_agent_yields_empty() {
871 let cfg = ModelSwitchConfig::default();
873 assert!(resolve_model_refs(&profile(None, &[]), &cfg, None).is_empty());
874 }
875
876 #[test]
877 fn difficulty_picks_frontier_over_threshold() {
878 let r = RoutingConfig {
879 enabled: true,
880 cheap: Some("cheap".into()),
881 frontier: Some("frontier".into()),
882 threshold_input_tokens: Some(1000),
883 };
884 assert_eq!(choose_by_difficulty(1500, &r), Some("frontier".into()));
885 assert_eq!(choose_by_difficulty(500, &r), Some("cheap".into()));
886 let bad = RoutingConfig {
888 enabled: true,
889 cheap: Some("c".into()),
890 frontier: None,
891 threshold_input_tokens: None,
892 };
893 assert_eq!(choose_by_difficulty(9999, &bad), None);
894 }
895
896 #[test]
897 fn pick_cheap_model_lowest_cost_chat_excluding_primary() {
898 let mut reg = ModelRegistry::default();
899 let mk = |cost: f64, caps: &[&str]| ModelEntry {
900 provider: "x".into(),
901 model: "m".into(),
902 capabilities: caps.iter().map(|s| s.to_string()).collect(),
903 cost_per_1k_tokens: Some(cost),
904 ..Default::default()
905 };
906 reg.models.insert("frontier".into(), mk(0.01, &["chat"]));
907 reg.models.insert("cheap".into(), mk(0.0001, &["chat"]));
908 reg.models
909 .insert("embed".into(), mk(0.00001, &["embedding"])); assert_eq!(
912 pick_cheap_model(®, Some("cheap"), &[]),
913 Some("frontier".into())
914 ); assert_eq!(pick_cheap_model(®, None, &[]), Some("cheap".into()));
916 let mut empty = ModelRegistry::default();
918 empty.models.insert("e".into(), mk(0.0, &["embedding"]));
919 assert_eq!(pick_cheap_model(&empty, None, &[]), None);
920 }
921
922 #[test]
923 fn satisfies_is_permissive_at_baseline_and_fail_closed_above_it() {
924 let mk = |caps: &[&str]| ModelEntry {
925 provider: "x".into(),
926 model: "m".into(),
927 capabilities: caps.iter().map(|s| s.to_string()).collect(),
928 ..Default::default()
929 };
930 assert!(satisfies(&mk(&[]), &[]));
932 assert!(satisfies(&mk(&["chat"]), &[]));
933 assert!(!satisfies(&mk(&["embedding"]), &[]));
934 assert!(!satisfies(&mk(&[]), &[Requirement::Vision]));
936 assert!(!satisfies(&mk(&["chat"]), &[Requirement::Vision]));
937 assert!(satisfies(&mk(&["chat", "vision"]), &[Requirement::Vision]));
938 assert!(!satisfies(&mk(&["chat", "vision"]), &[Requirement::Tools]));
939 assert!(satisfies(
940 &mk(&["chat", "vision", "tools"]),
941 &[Requirement::Vision, Requirement::Tools]
942 ));
943 }
944
945 #[test]
952 fn undeclared_capabilities_pass_tools_but_never_vision() {
953 let mk = |caps: &[&str]| ModelEntry {
954 provider: "x".into(),
955 model: "m".into(),
956 capabilities: caps.iter().map(|s| s.to_string()).collect(),
957 ..Default::default()
958 };
959 assert!(satisfies(&mk(&[]), &[Requirement::Tools]));
961 assert!(!satisfies(&mk(&[]), &[Requirement::Vision]));
962 assert!(!satisfies(
963 &mk(&[]),
964 &[Requirement::Vision, Requirement::Tools]
965 ));
966 assert!(!satisfies(&mk(&["chat"]), &[Requirement::Tools]));
968 assert!(satisfies(&mk(&["chat", "tools"]), &[Requirement::Tools]));
969 }
970
971 #[test]
974 fn pick_cheap_model_declines_when_no_entry_declares_the_requirement() {
975 let mk = |cost: f64, caps: &[&str]| ModelEntry {
976 provider: "x".into(),
977 model: "m".into(),
978 capabilities: caps.iter().map(|s| s.to_string()).collect(),
979 cost_per_1k_tokens: Some(cost),
980 ..Default::default()
981 };
982 let mut reg = ModelRegistry::default();
983 reg.models
984 .insert("cheap_text".into(), mk(0.0001, &["chat"]));
985 reg.models.insert("legacy".into(), mk(0.0002, &[]));
986 reg.models
987 .insert("frontier".into(), mk(0.01, &["chat", "vision"]));
988 assert_eq!(pick_cheap_model(®, None, &[]), Some("cheap_text".into()));
990 assert_eq!(
992 pick_cheap_model(®, None, &[Requirement::Vision]),
993 Some("frontier".into())
994 );
995 let mut blind = ModelRegistry::default();
997 blind
998 .models
999 .insert("cheap_text".into(), mk(0.0001, &["chat"]));
1000 blind.models.insert("legacy".into(), mk(0.0002, &[]));
1001 assert_eq!(pick_cheap_model(&blind, None, &[Requirement::Vision]), None);
1002 }
1003
1004 #[test]
1008 fn priced_at_stamps_only_priced_entries() {
1009 let now = chrono::Utc::now();
1010
1011 let mut unpriced = ModelEntry {
1012 provider: "openai".into(),
1013 model: "local-thing".into(),
1014 ..Default::default()
1015 };
1016 unpriced.stamp_priced_at(now);
1017 assert_eq!(unpriced.priced_at, None);
1018 assert_eq!(unpriced.price_age(now), None);
1019
1020 let mut priced = ModelEntry {
1021 output_cost_per_1k: Some(0.025),
1022 ..unpriced.clone()
1023 };
1024 priced.stamp_priced_at(now);
1025 assert_eq!(priced.priced_at, Some(now));
1026
1027 let mut legacy = ModelEntry {
1029 cost_per_1k_tokens: Some(0.01),
1030 ..unpriced.clone()
1031 };
1032 legacy.stamp_priced_at(now);
1033 assert!(legacy.priced_at.is_some());
1034
1035 let earlier = now - chrono::TimeDelta::days(30);
1037 priced.stamp_priced_at(earlier);
1038 assert_eq!(priced.priced_at, Some(now));
1039 }
1040
1041 #[test]
1044 fn registry_without_priced_at_still_loads_and_reports_unknown_age() {
1045 let yaml = r#"
1046schema_version: 1
1047models:
1048 opus:
1049 provider: anthropic
1050 model: claude-opus-5
1051 input_cost_per_1k: 0.005
1052 output_cost_per_1k: 0.025
1053"#;
1054 let reg: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
1055 let e = ®.models["opus"];
1056 assert_eq!(e.priced_at, None);
1057 assert_eq!(e.price_age(chrono::Utc::now()), None);
1058 let out = serde_yaml_ng::to_string(®).unwrap();
1060 assert!(!out.contains("priced_at"), "{out}");
1061 }
1062
1063 #[test]
1066 fn pick_cheap_model_sees_split_cost_entries() {
1067 let mut reg = ModelRegistry::default();
1068 let split = |input: f64, output: f64| ModelEntry {
1069 provider: "x".into(),
1070 model: "m".into(),
1071 capabilities: vec!["chat".into()],
1072 input_cost_per_1k: Some(input),
1073 output_cost_per_1k: Some(output),
1074 ..Default::default()
1075 };
1076 reg.models.insert("dear".into(), split(0.005, 0.025));
1077 reg.models.insert("cheap".into(), split(0.0001, 0.0004));
1078 assert_eq!(pick_cheap_model(®, None, &[]), Some("cheap".into()));
1079
1080 let mut input_only = ModelRegistry::default();
1082 input_only.models.insert(
1083 "in".into(),
1084 ModelEntry {
1085 provider: "x".into(),
1086 model: "m".into(),
1087 capabilities: vec!["chat".into()],
1088 input_cost_per_1k: Some(0.002),
1089 ..Default::default()
1090 },
1091 );
1092 assert_eq!(pick_cheap_model(&input_only, None, &[]), Some("in".into()));
1093 }
1094}