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 fn pick_cheap_model(reg: &ModelRegistry, exclude: Option<&str>) -> Option<String> {
298 reg.models
299 .iter()
300 .filter(|(k, _)| exclude != Some(k.as_str()))
301 .filter(|(_, e)| e.capabilities.is_empty() || e.capabilities.iter().any(|c| c == "chat"))
302 .filter_map(|(k, e)| {
303 let (input, output) = e.effective_costs();
307 output.or(input).map(|c| (c, k.clone()))
308 })
309 .min_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal))
310 .map(|(_, k)| k)
311}
312
313#[cfg(test)]
314mod tests {
315
316 #[test]
317 fn vendor_candidates_prefer_the_recorded_vendor_then_the_host_then_provider() {
318 let e = ModelEntry {
320 provider: "openai".into(),
321 vendor: Some("deepseek".into()),
322 base_url: Some("https://api.deepseek.com/v1".into()),
323 ..Default::default()
324 };
325 assert_eq!(e.vendor_candidates(), vec!["deepseek", "openai"]);
326
327 let legacy = ModelEntry {
330 provider: "openai".into(),
331 base_url: Some("https://api.deepseek.com/v1".into()),
332 ..Default::default()
333 };
334 assert_eq!(legacy.vendor_candidates(), vec!["deepseek", "openai"]);
335
336 let bare = ModelEntry {
338 provider: "anthropic".into(),
339 ..Default::default()
340 };
341 assert_eq!(bare.vendor_candidates(), vec!["anthropic"]);
342
343 let same = ModelEntry {
345 provider: "openai".into(),
346 base_url: Some("https://api.openai.com/v1".into()),
347 ..Default::default()
348 };
349 assert_eq!(same.vendor_candidates(), vec!["openai"]);
350 }
351
352 #[test]
353 fn vendor_is_omitted_from_yaml_when_absent_and_round_trips_when_set() {
354 let bare = ModelEntry {
355 provider: "anthropic".into(),
356 model: "claude-opus-5".into(),
357 ..Default::default()
358 };
359 let y = serde_yaml_ng::to_string(&bare).unwrap();
360 assert!(!y.contains("vendor"), "{y}");
361
362 let tagged = ModelEntry {
363 provider: "openai".into(),
364 vendor: Some("groq".into()),
365 model: "llama-3.3".into(),
366 ..Default::default()
367 };
368 let y = serde_yaml_ng::to_string(&tagged).unwrap();
369 let back: ModelEntry = serde_yaml_ng::from_str(&y).unwrap();
370 assert_eq!(back.vendor.as_deref(), Some("groq"));
371 }
372 use super::*;
373
374 #[test]
375 fn parses_full_registry() {
376 let yaml = r#"
377schema_version: 1
378models:
379 anthropic_opus_4_7:
380 provider: anthropic
381 model: claude-opus-4-7
382 secret: env:ANTHROPIC_API_KEY
383 capabilities: [chat, tools]
384 ollama_llama3:
385 provider: ollama
386 model: llama3.2:3b
387 base_url: http://127.0.0.1:11434
388"#;
389 let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
390 assert_eq!(r.schema_version, 1);
391 assert_eq!(r.models.len(), 2);
392 let opus = r.models.get("anthropic_opus_4_7").unwrap();
393 assert_eq!(opus.provider, "anthropic");
394 assert_eq!(
395 opus.secret,
396 Some(SecretRef::Env("ANTHROPIC_API_KEY".into()))
397 );
398 assert!(r.models["ollama_llama3"].secret.is_none());
399 }
400
401 #[test]
402 fn round_trip_preserves_shape() {
403 let mut r = ModelRegistry::default();
404 r.models.insert(
405 "foo".into(),
406 ModelEntry {
407 provider: "anthropic".into(),
408 model: "claude-opus-4-7".into(),
409 base_url: None,
410 secret: Some(SecretRef::Keychain {
411 service: "mur".into(),
412 account: "anthropic".into(),
413 }),
414 capabilities: vec!["chat".into()],
415 params: serde_json::Value::Null,
416 tier: None,
417 cost_per_1k_tokens: None,
418 input_cost_per_1k: None,
419 output_cost_per_1k: None,
420 context_window: None,
421 priced_at: None,
422 ..Default::default()
423 },
424 );
425 let s = serde_yaml_ng::to_string(&r).unwrap();
426 let parsed: ModelRegistry = serde_yaml_ng::from_str(&s).unwrap();
427 assert_eq!(r, parsed);
428 }
429
430 #[test]
431 fn rejects_unknown_secret_scheme() {
432 let yaml = r#"
433schema_version: 1
434models:
435 bad:
436 provider: x
437 model: y
438 secret: bogus:value
439"#;
440 let r: Result<ModelRegistry, _> = serde_yaml_ng::from_str(yaml);
441 assert!(r.is_err(), "should reject unknown scheme");
442 }
443
444 #[test]
445 fn test_registry_roundtrip_with_roles() {
446 let yaml = r#"
447schema_version: 1
448models:
449 haiku:
450 provider: anthropic
451 model: claude-haiku-4-5
452roles:
453 reflector:
454 primary: haiku
455 fallback: null
456 cost_budget_per_day_usd: 0.5
457"#;
458 let reg: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
459 assert_eq!(reg.roles["reflector"].primary, "haiku");
460 let back = serde_yaml_ng::to_string(®).unwrap();
461 let reg2: ModelRegistry = serde_yaml_ng::from_str(&back).unwrap();
462 assert_eq!(reg, reg2);
463 }
464
465 #[test]
466 fn test_resolve_role_primary() {
467 let mut reg = ModelRegistry::default();
468 reg.models.insert(
469 "haiku".into(),
470 ModelEntry {
471 provider: "anthropic".into(),
472 model: "claude-haiku-4-5".into(),
473 base_url: None,
474 secret: None,
475 capabilities: vec![],
476 params: serde_json::Value::Null,
477 tier: None,
478 cost_per_1k_tokens: None,
479 input_cost_per_1k: None,
480 output_cost_per_1k: None,
481 context_window: None,
482 priced_at: None,
483 ..Default::default()
484 },
485 );
486 reg.roles.insert(
487 "reflector".into(),
488 RoleEntry {
489 primary: "haiku".into(),
490 fallback: None,
491 ..Default::default()
492 },
493 );
494 assert_eq!(reg.resolve_role("reflector"), Some("haiku"));
495 }
496
497 #[test]
498 fn test_resolve_role_fallback() {
499 let mut reg = ModelRegistry::default();
500 reg.models.insert(
501 "haiku".into(),
502 ModelEntry {
503 provider: "anthropic".into(),
504 model: "claude-haiku-4-5".into(),
505 base_url: None,
506 secret: None,
507 capabilities: vec![],
508 params: serde_json::Value::Null,
509 tier: None,
510 cost_per_1k_tokens: None,
511 input_cost_per_1k: None,
512 output_cost_per_1k: None,
513 context_window: None,
514 priced_at: None,
515 ..Default::default()
516 },
517 );
518 reg.roles.insert(
519 "reflector".into(),
520 RoleEntry {
521 primary: "nonexistent".into(),
522 fallback: Some("haiku".into()),
523 ..Default::default()
524 },
525 );
526 assert_eq!(reg.resolve_role("reflector"), Some("haiku"));
527 }
528
529 #[test]
530 fn test_resolve_role_none() {
531 let reg = ModelRegistry::default();
532 assert_eq!(reg.resolve_role("reflector"), None);
533 }
534
535 #[test]
536 fn model_entry_parses_tier_field() {
537 let yaml = r#"
538schema_version: 1
539models:
540 haiku:
541 provider: anthropic
542 model: claude-haiku-4-5
543 tier: local
544 opus:
545 provider: anthropic
546 model: claude-opus-4-7
547 tier: frontier
548 cost_per_1k_tokens: 0.015
549"#;
550 let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
551 assert_eq!(r.models["haiku"].tier, Some(RouteTier::Local));
552 assert_eq!(r.models["opus"].tier, Some(RouteTier::Frontier));
553 assert_eq!(r.models["opus"].cost_per_1k_tokens, Some(0.015));
554 let mut r2 = ModelRegistry::default();
556 r2.models.insert(
557 "x".into(),
558 ModelEntry {
559 provider: "ollama".into(),
560 model: "llama3".into(),
561 base_url: None,
562 secret: None,
563 capabilities: vec![],
564 params: serde_json::Value::Null,
565 tier: None,
566 cost_per_1k_tokens: None,
567 input_cost_per_1k: None,
568 output_cost_per_1k: None,
569 context_window: None,
570 priced_at: None,
571 ..Default::default()
572 },
573 );
574 let yaml = serde_yaml_ng::to_string(&r2).unwrap();
575 assert!(
576 !yaml.contains("tier:"),
577 "absent tier should not be serialized: {yaml}"
578 );
579 }
580
581 #[test]
582 fn role_entry_parses_route_policy() {
583 let yaml = r#"
584schema_version: 1
585models:
586 haiku:
587 provider: anthropic
588 model: claude-haiku-4-5
589 opus:
590 provider: anthropic
591 model: claude-opus-4-7
592roles:
593 dev:
594 primary: opus
595 route_policy: !force_frontier
596 model_id: opus
597 reflector:
598 primary: haiku
599 route_policy: prefer_local
600 curator:
601 primary: haiku
602 route_policy: force_local
603 chat:
604 primary: haiku
605"#;
606 let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
607 assert_eq!(
608 r.roles["dev"].route_policy,
609 Some(RoutePolicy::ForceFrontier {
610 model_id: "opus".into()
611 })
612 );
613 assert_eq!(
614 r.roles["reflector"].route_policy,
615 Some(RoutePolicy::PreferLocal)
616 );
617 assert_eq!(
618 r.roles["curator"].route_policy,
619 Some(RoutePolicy::ForceLocal)
620 );
621 assert_eq!(r.roles["chat"].route_policy, None);
622 }
623
624 #[test]
625 fn parses_split_cost_fields() {
626 let yaml = r#"
627schema_version: 1
628models:
629 opus:
630 provider: anthropic
631 model: claude-opus-4-8
632 input_cost_per_1k: 0.005
633 output_cost_per_1k: 0.025
634 context_window: 200000
635"#;
636 let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
637 let e = r.models.get("opus").unwrap();
638 assert_eq!(e.input_cost_per_1k, Some(0.005));
639 assert_eq!(e.output_cost_per_1k, Some(0.025));
640 assert_eq!(e.context_window, Some(200_000));
641 }
642
643 #[test]
644 fn default_model_entry_is_empty() {
645 let e = ModelEntry::default();
646 assert!(e.provider.is_empty());
647 assert_eq!(e.input_cost_per_1k, None);
648 assert_eq!(e.output_cost_per_1k, None);
649 assert_eq!(e.context_window, None);
650 }
651
652 #[test]
653 fn effective_costs_fallback_matrix() {
654 let mut e = ModelEntry {
656 cost_per_1k_tokens: Some(0.01),
657 ..Default::default()
658 };
659 assert_eq!(e.effective_costs(), (Some(0.01), Some(0.01)));
660
661 e = ModelEntry {
663 input_cost_per_1k: Some(0.005),
664 output_cost_per_1k: Some(0.025),
665 ..Default::default()
666 };
667 assert_eq!(e.effective_costs(), (Some(0.005), Some(0.025)));
668
669 e = ModelEntry {
671 cost_per_1k_tokens: Some(0.01),
672 input_cost_per_1k: Some(0.005),
673 output_cost_per_1k: Some(0.025),
674 ..Default::default()
675 };
676 assert_eq!(e.effective_costs(), (Some(0.005), Some(0.025)));
677
678 e = ModelEntry::default();
680 assert_eq!(e.effective_costs(), (None, None));
681 }
682}
683
684#[cfg(test)]
685mod io_tests {
686 use super::*;
687 use tempfile::tempdir;
688
689 #[test]
690 fn load_returns_empty_when_file_missing() {
691 let dir = tempdir().unwrap();
692 let r = ModelRegistry::load_from(&dir.path().join("nope.yaml")).unwrap();
693 assert_eq!(r.models.len(), 0);
694 assert_eq!(r.schema_version, 1);
695 }
696
697 #[test]
698 fn save_then_load_round_trips() {
699 let dir = tempdir().unwrap();
700 let p = dir.path().join("models.yaml");
701 let mut r = ModelRegistry::default();
702 r.models.insert(
703 "x".into(),
704 ModelEntry {
705 provider: "ollama".into(),
706 model: "llama3.2:3b".into(),
707 base_url: None,
708 secret: None,
709 capabilities: vec![],
710 params: serde_json::Value::Null,
711 tier: None,
712 cost_per_1k_tokens: None,
713 input_cost_per_1k: None,
714 output_cost_per_1k: None,
715 context_window: None,
716 priced_at: None,
717 ..Default::default()
718 },
719 );
720 r.save_to(&p).unwrap();
721 let r2 = ModelRegistry::load_from(&p).unwrap();
722 assert_eq!(r, r2);
723 }
724
725 #[test]
726 fn save_uses_atomic_rename() {
727 let dir = tempdir().unwrap();
728 let p = dir.path().join("models.yaml");
729 ModelRegistry::default().save_to(&p).unwrap();
730 let temp = dir.path().join("models.yaml.tmp");
731 assert!(!temp.exists(), "atomic temp left behind");
732 }
733}
734
735#[cfg(test)]
736mod switch_tests {
737 use super::*;
738 use crate::agent::AgentProfile;
739 use crate::config::{ModelSwitchConfig, RoutingConfig};
740
741 fn profile(model_ref: Option<&str>, chain: &[&str]) -> AgentProfile {
742 let mut p = AgentProfile::default_for_tests();
743 p.model_ref = model_ref.map(|s| s.to_string());
744 p.fallback_chain = chain.iter().map(|s| s.to_string()).collect();
745 p
746 }
747
748 #[test]
749 fn per_agent_primary_and_chain_win_over_global() {
750 let cfg = ModelSwitchConfig {
751 default: Some("global_default".into()),
752 fallback_chain: vec!["g1".into(), "g2".into()],
753 ..Default::default()
754 };
755 let p = profile(Some("agent_primary"), &["agent_primary", "agent_fb"]);
756 assert_eq!(
758 resolve_model_refs(&p, &cfg, None),
759 vec!["agent_primary", "agent_fb"]
760 );
761 }
762
763 #[test]
764 fn falls_back_to_global_default_and_chain() {
765 let cfg = ModelSwitchConfig {
766 default: Some("global_default".into()),
767 fallback_chain: vec!["g1".into(), "global_default".into()],
768 ..Default::default()
769 };
770 let p = profile(None, &[]); assert_eq!(
773 resolve_model_refs(&p, &cfg, None),
774 vec!["global_default", "g1"]
775 );
776 }
777
778 #[test]
779 fn routed_primary_overrides_model_ref() {
780 let cfg = ModelSwitchConfig {
781 fallback_chain: vec!["g1".into()],
782 ..Default::default()
783 };
784 let p = profile(Some("agent_primary"), &[]);
785 assert_eq!(
786 resolve_model_refs(&p, &cfg, Some("frontier".into())),
787 vec!["frontier", "g1"]
788 );
789 }
790
791 #[test]
792 fn no_config_no_agent_yields_empty() {
793 let cfg = ModelSwitchConfig::default();
795 assert!(resolve_model_refs(&profile(None, &[]), &cfg, None).is_empty());
796 }
797
798 #[test]
799 fn difficulty_picks_frontier_over_threshold() {
800 let r = RoutingConfig {
801 enabled: true,
802 cheap: Some("cheap".into()),
803 frontier: Some("frontier".into()),
804 threshold_input_tokens: Some(1000),
805 ..Default::default()
806 };
807 assert_eq!(choose_by_difficulty(1500, &r), Some("frontier".into()));
808 assert_eq!(choose_by_difficulty(500, &r), Some("cheap".into()));
809 let bad = RoutingConfig {
811 enabled: true,
812 cheap: Some("c".into()),
813 frontier: None,
814 threshold_input_tokens: None,
815 ..Default::default()
816 };
817 assert_eq!(choose_by_difficulty(9999, &bad), None);
818 }
819
820 #[test]
821 fn pick_cheap_model_lowest_cost_chat_excluding_primary() {
822 let mut reg = ModelRegistry::default();
823 let mk = |cost: f64, caps: &[&str]| ModelEntry {
824 provider: "x".into(),
825 model: "m".into(),
826 capabilities: caps.iter().map(|s| s.to_string()).collect(),
827 cost_per_1k_tokens: Some(cost),
828 ..Default::default()
829 };
830 reg.models.insert("frontier".into(), mk(0.01, &["chat"]));
831 reg.models.insert("cheap".into(), mk(0.0001, &["chat"]));
832 reg.models
833 .insert("embed".into(), mk(0.00001, &["embedding"])); assert_eq!(
836 pick_cheap_model(®, Some("cheap")),
837 Some("frontier".into())
838 ); assert_eq!(pick_cheap_model(®, None), Some("cheap".into()));
840 let mut empty = ModelRegistry::default();
842 empty.models.insert("e".into(), mk(0.0, &["embedding"]));
843 assert_eq!(pick_cheap_model(&empty, None), None);
844 }
845
846 #[test]
850 fn priced_at_stamps_only_priced_entries() {
851 let now = chrono::Utc::now();
852
853 let mut unpriced = ModelEntry {
854 provider: "openai".into(),
855 model: "local-thing".into(),
856 ..Default::default()
857 };
858 unpriced.stamp_priced_at(now);
859 assert_eq!(unpriced.priced_at, None);
860 assert_eq!(unpriced.price_age(now), None);
861
862 let mut priced = ModelEntry {
863 output_cost_per_1k: Some(0.025),
864 ..unpriced.clone()
865 };
866 priced.stamp_priced_at(now);
867 assert_eq!(priced.priced_at, Some(now));
868
869 let mut legacy = ModelEntry {
871 cost_per_1k_tokens: Some(0.01),
872 ..unpriced.clone()
873 };
874 legacy.stamp_priced_at(now);
875 assert!(legacy.priced_at.is_some());
876
877 let earlier = now - chrono::TimeDelta::days(30);
879 priced.stamp_priced_at(earlier);
880 assert_eq!(priced.priced_at, Some(now));
881 }
882
883 #[test]
886 fn registry_without_priced_at_still_loads_and_reports_unknown_age() {
887 let yaml = r#"
888schema_version: 1
889models:
890 opus:
891 provider: anthropic
892 model: claude-opus-5
893 input_cost_per_1k: 0.005
894 output_cost_per_1k: 0.025
895"#;
896 let reg: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
897 let e = ®.models["opus"];
898 assert_eq!(e.priced_at, None);
899 assert_eq!(e.price_age(chrono::Utc::now()), None);
900 let out = serde_yaml_ng::to_string(®).unwrap();
902 assert!(!out.contains("priced_at"), "{out}");
903 }
904
905 #[test]
908 fn pick_cheap_model_sees_split_cost_entries() {
909 let mut reg = ModelRegistry::default();
910 let split = |input: f64, output: f64| ModelEntry {
911 provider: "x".into(),
912 model: "m".into(),
913 capabilities: vec!["chat".into()],
914 input_cost_per_1k: Some(input),
915 output_cost_per_1k: Some(output),
916 ..Default::default()
917 };
918 reg.models.insert("dear".into(), split(0.005, 0.025));
919 reg.models.insert("cheap".into(), split(0.0001, 0.0004));
920 assert_eq!(pick_cheap_model(®, None), Some("cheap".into()));
921
922 let mut input_only = ModelRegistry::default();
924 input_only.models.insert(
925 "in".into(),
926 ModelEntry {
927 provider: "x".into(),
928 model: "m".into(),
929 capabilities: vec!["chat".into()],
930 input_cost_per_1k: Some(0.002),
931 ..Default::default()
932 },
933 );
934 assert_eq!(pick_cheap_model(&input_only, None), Some("in".into()));
935 }
936}