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, Copy, Serialize, Deserialize, PartialEq, Eq)]
28#[serde(rename_all = "snake_case")]
29pub enum BillingMode {
30 Subscription,
32 UsageBilled,
34 Local,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
39pub struct ModelEntry {
40 #[serde(default)]
41 pub provider: String,
42 #[serde(default, skip_serializing_if = "Option::is_none")]
54 pub vendor: Option<String>,
55 #[serde(default)]
56 pub model: String,
57 #[serde(default, skip_serializing_if = "Option::is_none")]
58 pub base_url: Option<String>,
59 #[serde(default, skip_serializing_if = "Option::is_none")]
60 pub secret: Option<SecretRef>,
61 #[serde(default, skip_serializing_if = "Vec::is_empty")]
62 pub capabilities: Vec<String>,
63 #[serde(default, skip_serializing_if = "serde_json::Value::is_null")]
64 pub params: serde_json::Value,
65 #[serde(default, skip_serializing_if = "Option::is_none")]
68 pub tier: Option<RouteTier>,
69 #[serde(default, skip_serializing_if = "Option::is_none")]
72 pub cost_per_1k_tokens: Option<f64>,
73 #[serde(default, skip_serializing_if = "Option::is_none")]
76 pub input_cost_per_1k: Option<f64>,
77 #[serde(default, skip_serializing_if = "Option::is_none")]
80 pub output_cost_per_1k: Option<f64>,
81 #[serde(default, skip_serializing_if = "Option::is_none")]
83 pub context_window: Option<u64>,
84 #[serde(default, skip_serializing_if = "Option::is_none")]
91 pub priced_at: Option<chrono::DateTime<chrono::Utc>>,
92 #[serde(default, skip_serializing_if = "Option::is_none")]
94 pub billing: Option<BillingMode>,
95 #[serde(default, skip_serializing_if = "Option::is_none")]
99 pub catalog_verified: Option<bool>,
100}
101
102fn vendor_label_of_url(base_url: Option<&str>) -> Option<String> {
107 let host = base_url?
108 .split("//")
109 .nth(1)
110 .unwrap_or(base_url?)
111 .split(['/', ':'])
112 .next()
113 .unwrap_or("");
114 let label = host.strip_prefix("api.").unwrap_or(host);
115 let first = label.split('.').next().unwrap_or("");
116 (!first.is_empty()).then(|| first.to_string())
117}
118
119const LOCAL_MODEL_PROVIDERS: &[&str] = &[
120 "ollama",
121 "mlx",
122 "llamacpp",
123 "llama_cpp",
124 "localai",
125 "lmstudio",
126 "local",
127];
128
129pub fn provider_is_local(provider: &str) -> bool {
131 LOCAL_MODEL_PROVIDERS.contains(&provider.to_ascii_lowercase().as_str())
132}
133
134pub fn inferred_billing_for_provider(provider: &str) -> BillingMode {
136 if provider_is_local(provider) {
137 BillingMode::Local
138 } else if matches!(provider.to_ascii_lowercase().as_str(), "codex" | "claude") {
139 BillingMode::Subscription
140 } else {
141 BillingMode::UsageBilled
142 }
143}
144
145impl ModelEntry {
146 pub fn billing_or_inferred(&self) -> BillingMode {
155 self.billing
156 .unwrap_or_else(|| inferred_billing_for_provider(&self.provider))
157 }
158
159 pub fn effective_route_tier(&self) -> RouteTier {
161 self.tier.unwrap_or_else(|| {
162 if provider_is_local(&self.provider) {
163 RouteTier::Local
164 } else {
165 RouteTier::Frontier
166 }
167 })
168 }
169 pub fn effective_costs(&self) -> (Option<f64>, Option<f64>) {
174 let output = self.output_cost_per_1k.or(self.cost_per_1k_tokens);
175 let input = self.input_cost_per_1k.or(self.cost_per_1k_tokens);
176 (input, output)
177 }
178
179 pub fn projected_cost(&self, input_tokens: u64, output_tokens: u64) -> Option<f64> {
182 let (input, output) = self.effective_costs();
183 let fallback = input.or(output)?;
184 let input = input.unwrap_or(fallback);
185 let output = output.unwrap_or(fallback);
186 Some(input_tokens as f64 / 1_000.0 * input + output_tokens as f64 / 1_000.0 * output)
187 }
188
189 pub fn vendor_candidates(&self) -> Vec<String> {
200 let mut out: Vec<String> = Vec::with_capacity(3);
201 let mut push = |v: &str| {
202 if !v.is_empty() && !out.iter().any(|e| e == v) {
203 out.push(v.to_string());
204 }
205 };
206 if let Some(v) = self.vendor.as_deref() {
207 push(v);
208 }
209 if let Some(label) = vendor_label_of_url(self.base_url.as_deref()) {
210 push(&label);
211 }
212 push(&self.provider);
213 out
214 }
215
216 pub fn is_priced(&self) -> bool {
218 let (input, output) = self.effective_costs();
219 input.is_some() || output.is_some()
220 }
221
222 pub fn stamp_priced_at(&mut self, now: chrono::DateTime<chrono::Utc>) {
226 if self.is_priced() && self.priced_at.is_none_or(|prev| prev < now) {
227 self.priced_at = Some(now);
228 }
229 }
230
231 pub fn price_age(&self, now: chrono::DateTime<chrono::Utc>) -> Option<chrono::TimeDelta> {
233 self.priced_at.map(|at| now - at)
234 }
235}
236
237#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
238pub struct RoleEntry {
239 pub primary: String,
241 #[serde(default, skip_serializing_if = "Option::is_none")]
243 pub fallback: Option<String>,
244 #[serde(default, skip_serializing_if = "Option::is_none")]
246 pub cost_budget_per_day_usd: Option<f64>,
247 #[serde(default)]
249 pub privacy_local_only: bool,
250 #[serde(default, skip_serializing_if = "Option::is_none")]
253 pub route_policy: Option<RoutePolicy>,
254}
255
256#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
257pub struct ModelRegistry {
258 pub schema_version: u32,
259 #[serde(default)]
260 pub models: BTreeMap<String, ModelEntry>,
261 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
262 pub roles: BTreeMap<String, RoleEntry>,
263}
264
265impl Default for ModelRegistry {
266 fn default() -> Self {
267 Self {
268 schema_version: 1,
269 models: BTreeMap::new(),
270 roles: BTreeMap::new(),
271 }
272 }
273}
274
275impl ModelRegistry {
276 pub fn load_from(path: &Path) -> anyhow::Result<Self> {
277 if !path.exists() {
278 return Ok(Self::default());
279 }
280 let body = std::fs::read_to_string(path)?;
281 if body.trim().is_empty() {
282 return Ok(Self::default());
283 }
284 Ok(serde_yaml_ng::from_str(&body)?)
285 }
286
287 pub fn save_to(&self, path: &Path) -> anyhow::Result<()> {
288 if let Some(parent) = path.parent() {
289 std::fs::create_dir_all(parent)?;
290 }
291 let body = serde_yaml_ng::to_string(self)?;
292 let tmp = path.with_extension("yaml.tmp");
293 std::fs::write(&tmp, body)?;
294 std::fs::rename(&tmp, path)?;
295 Ok(())
296 }
297
298 pub fn default_path() -> anyhow::Result<PathBuf> {
299 if let Ok(p) = std::env::var("MUR_HOME")
302 && !p.is_empty()
303 {
304 return Ok(PathBuf::from(p).join("models.yaml"));
305 }
306 let home = dirs::home_dir().ok_or_else(|| anyhow::anyhow!("no home dir"))?;
307 Ok(home.join(".mur/models.yaml"))
308 }
309
310 pub fn resolve_role(&self, role: &str) -> Option<&str> {
313 let entry = self.roles.get(role)?;
314 if self.models.contains_key(&entry.primary) {
315 return Some(&entry.primary);
316 }
317 if let Some(fb) = &entry.fallback
319 && self.models.contains_key(fb)
320 {
321 return Some(fb);
322 }
323 None
325 }
326}
327
328use crate::agent::AgentProfile;
329use crate::config::{DEFAULT_ROUTING_THRESHOLD, ModelSwitchConfig, RoutingConfig};
330
331pub fn resolve_model_refs(
336 profile: &AgentProfile,
337 cfg: &ModelSwitchConfig,
338 routed_primary: Option<String>,
339) -> Vec<String> {
340 let primary = routed_primary
341 .or_else(|| profile.model_ref.clone())
342 .or_else(|| cfg.default.clone());
343 let chain = if !profile.fallback_chain.is_empty() {
344 profile.fallback_chain.clone()
345 } else {
346 cfg.fallback_chain.clone()
347 };
348 let mut out: Vec<String> = Vec::new();
349 if let Some(p) = primary {
350 out.push(p);
351 }
352 for r in chain {
353 if !out.contains(&r) {
354 out.push(r);
355 }
356 }
357 out
358}
359
360pub fn choose_by_difficulty(est_input_tokens: u32, r: &RoutingConfig) -> Option<String> {
364 let threshold = r
365 .threshold_input_tokens
366 .unwrap_or(DEFAULT_ROUTING_THRESHOLD);
367 match (r.cheap.as_ref(), r.frontier.as_ref()) {
368 (Some(cheap), Some(frontier)) => Some(if est_input_tokens > threshold {
369 frontier.clone()
370 } else {
371 cheap.clone()
372 }),
373 _ => None,
374 }
375}
376
377pub const CAP_CHAT: &str = "chat";
381pub const CAP_TOOLS: &str = "tools";
382pub const CAP_VISION: &str = "vision";
383
384#[derive(Debug, Clone, Copy, PartialEq, Eq)]
388pub enum Requirement {
389 Vision,
391 Tools,
393}
394
395impl Requirement {
396 pub fn capability(self) -> &'static str {
398 match self {
399 Requirement::Vision => CAP_VISION,
400 Requirement::Tools => CAP_TOOLS,
401 }
402 }
403
404 fn permitted_when_undeclared(self) -> bool {
424 match self {
425 Requirement::Vision => false,
426 Requirement::Tools => true,
427 }
428 }
429}
430
431pub fn satisfies(e: &ModelEntry, reqs: &[Requirement]) -> bool {
439 let chat_capable = e.capabilities.is_empty() || e.capabilities.iter().any(|c| c == CAP_CHAT);
440 if !chat_capable {
441 return false;
442 }
443 reqs.iter().all(|r| {
444 if e.capabilities.is_empty() {
445 r.permitted_when_undeclared()
446 } else {
447 e.capabilities.iter().any(|c| c == r.capability())
448 }
449 })
450}
451
452pub fn pick_cheap_model(
456 reg: &ModelRegistry,
457 exclude: Option<&str>,
458 reqs: &[Requirement],
459) -> Option<String> {
460 reg.models
461 .iter()
462 .filter(|(k, _)| exclude != Some(k.as_str()))
463 .filter(|(_, e)| satisfies(e, reqs))
464 .filter_map(|(k, e)| e.projected_cost(4_000, 1_000).map(|c| (c, k.clone())))
465 .min_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal))
466 .map(|(_, k)| k)
467}
468
469#[cfg(test)]
470mod tests {
471
472 #[test]
473 fn vendor_candidates_prefer_the_recorded_vendor_then_the_host_then_provider() {
474 let e = ModelEntry {
476 provider: "openai".into(),
477 vendor: Some("deepseek".into()),
478 base_url: Some("https://api.deepseek.com/v1".into()),
479 ..Default::default()
480 };
481 assert_eq!(e.vendor_candidates(), vec!["deepseek", "openai"]);
482
483 let legacy = ModelEntry {
486 provider: "openai".into(),
487 base_url: Some("https://api.deepseek.com/v1".into()),
488 ..Default::default()
489 };
490 assert_eq!(legacy.vendor_candidates(), vec!["deepseek", "openai"]);
491
492 let bare = ModelEntry {
494 provider: "anthropic".into(),
495 ..Default::default()
496 };
497 assert_eq!(bare.vendor_candidates(), vec!["anthropic"]);
498
499 let same = ModelEntry {
501 provider: "openai".into(),
502 base_url: Some("https://api.openai.com/v1".into()),
503 ..Default::default()
504 };
505 assert_eq!(same.vendor_candidates(), vec!["openai"]);
506 }
507
508 #[test]
509 fn vendor_is_omitted_from_yaml_when_absent_and_round_trips_when_set() {
510 let bare = ModelEntry {
511 provider: "anthropic".into(),
512 model: "claude-opus-5".into(),
513 ..Default::default()
514 };
515 let y = serde_yaml_ng::to_string(&bare).unwrap();
516 assert!(!y.contains("vendor"), "{y}");
517
518 let tagged = ModelEntry {
519 provider: "openai".into(),
520 vendor: Some("groq".into()),
521 model: "llama-3.3".into(),
522 ..Default::default()
523 };
524 let y = serde_yaml_ng::to_string(&tagged).unwrap();
525 let back: ModelEntry = serde_yaml_ng::from_str(&y).unwrap();
526 assert_eq!(back.vendor.as_deref(), Some("groq"));
527 }
528 use super::*;
529
530 #[test]
531 fn parses_full_registry() {
532 let yaml = r#"
533schema_version: 1
534models:
535 anthropic_opus_4_7:
536 provider: anthropic
537 model: claude-opus-4-7
538 secret: env:ANTHROPIC_API_KEY
539 capabilities: [chat, tools]
540 ollama_llama3:
541 provider: ollama
542 model: llama3.2:3b
543 base_url: http://127.0.0.1:11434
544"#;
545 let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
546 assert_eq!(r.schema_version, 1);
547 assert_eq!(r.models.len(), 2);
548 let opus = r.models.get("anthropic_opus_4_7").unwrap();
549 assert_eq!(opus.provider, "anthropic");
550 assert_eq!(
551 opus.secret,
552 Some(SecretRef::Env("ANTHROPIC_API_KEY".into()))
553 );
554 assert!(r.models["ollama_llama3"].secret.is_none());
555 }
556
557 #[test]
558 fn round_trip_preserves_shape() {
559 let mut r = ModelRegistry::default();
560 r.models.insert(
561 "foo".into(),
562 ModelEntry {
563 provider: "anthropic".into(),
564 model: "claude-opus-4-7".into(),
565 base_url: None,
566 secret: Some(SecretRef::Keychain {
567 service: "mur".into(),
568 account: "anthropic".into(),
569 }),
570 capabilities: vec!["chat".into()],
571 params: serde_json::Value::Null,
572 tier: None,
573 cost_per_1k_tokens: None,
574 input_cost_per_1k: None,
575 output_cost_per_1k: None,
576 context_window: None,
577 priced_at: None,
578 ..Default::default()
579 },
580 );
581 let s = serde_yaml_ng::to_string(&r).unwrap();
582 let parsed: ModelRegistry = serde_yaml_ng::from_str(&s).unwrap();
583 assert_eq!(r, parsed);
584 }
585
586 #[test]
587 fn rejects_unknown_secret_scheme() {
588 let yaml = r#"
589schema_version: 1
590models:
591 bad:
592 provider: x
593 model: y
594 secret: bogus:value
595"#;
596 let r: Result<ModelRegistry, _> = serde_yaml_ng::from_str(yaml);
597 assert!(r.is_err(), "should reject unknown scheme");
598 }
599
600 #[test]
601 fn test_registry_roundtrip_with_roles() {
602 let yaml = r#"
603schema_version: 1
604models:
605 haiku:
606 provider: anthropic
607 model: claude-haiku-4-5
608roles:
609 reflector:
610 primary: haiku
611 fallback: null
612 cost_budget_per_day_usd: 0.5
613"#;
614 let reg: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
615 assert_eq!(reg.roles["reflector"].primary, "haiku");
616 let back = serde_yaml_ng::to_string(®).unwrap();
617 let reg2: ModelRegistry = serde_yaml_ng::from_str(&back).unwrap();
618 assert_eq!(reg, reg2);
619 }
620
621 #[test]
622 fn test_resolve_role_primary() {
623 let mut reg = ModelRegistry::default();
624 reg.models.insert(
625 "haiku".into(),
626 ModelEntry {
627 provider: "anthropic".into(),
628 model: "claude-haiku-4-5".into(),
629 base_url: None,
630 secret: None,
631 capabilities: vec![],
632 params: serde_json::Value::Null,
633 tier: None,
634 cost_per_1k_tokens: None,
635 input_cost_per_1k: None,
636 output_cost_per_1k: None,
637 context_window: None,
638 priced_at: None,
639 ..Default::default()
640 },
641 );
642 reg.roles.insert(
643 "reflector".into(),
644 RoleEntry {
645 primary: "haiku".into(),
646 fallback: None,
647 ..Default::default()
648 },
649 );
650 assert_eq!(reg.resolve_role("reflector"), Some("haiku"));
651 }
652
653 #[test]
654 fn test_resolve_role_fallback() {
655 let mut reg = ModelRegistry::default();
656 reg.models.insert(
657 "haiku".into(),
658 ModelEntry {
659 provider: "anthropic".into(),
660 model: "claude-haiku-4-5".into(),
661 base_url: None,
662 secret: None,
663 capabilities: vec![],
664 params: serde_json::Value::Null,
665 tier: None,
666 cost_per_1k_tokens: None,
667 input_cost_per_1k: None,
668 output_cost_per_1k: None,
669 context_window: None,
670 priced_at: None,
671 ..Default::default()
672 },
673 );
674 reg.roles.insert(
675 "reflector".into(),
676 RoleEntry {
677 primary: "nonexistent".into(),
678 fallback: Some("haiku".into()),
679 ..Default::default()
680 },
681 );
682 assert_eq!(reg.resolve_role("reflector"), Some("haiku"));
683 }
684
685 #[test]
686 fn test_resolve_role_none() {
687 let reg = ModelRegistry::default();
688 assert_eq!(reg.resolve_role("reflector"), None);
689 }
690
691 #[test]
692 fn model_entry_parses_tier_field() {
693 let yaml = r#"
694schema_version: 1
695models:
696 haiku:
697 provider: anthropic
698 model: claude-haiku-4-5
699 tier: local
700 opus:
701 provider: anthropic
702 model: claude-opus-4-7
703 tier: frontier
704 cost_per_1k_tokens: 0.015
705"#;
706 let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
707 assert_eq!(r.models["haiku"].tier, Some(RouteTier::Local));
708 assert_eq!(r.models["opus"].tier, Some(RouteTier::Frontier));
709 assert_eq!(r.models["opus"].cost_per_1k_tokens, Some(0.015));
710 let mut r2 = ModelRegistry::default();
712 r2.models.insert(
713 "x".into(),
714 ModelEntry {
715 provider: "ollama".into(),
716 model: "llama3".into(),
717 base_url: None,
718 secret: None,
719 capabilities: vec![],
720 params: serde_json::Value::Null,
721 tier: None,
722 cost_per_1k_tokens: None,
723 input_cost_per_1k: None,
724 output_cost_per_1k: None,
725 context_window: None,
726 priced_at: None,
727 ..Default::default()
728 },
729 );
730 let yaml = serde_yaml_ng::to_string(&r2).unwrap();
731 assert!(
732 !yaml.contains("tier:"),
733 "absent tier should not be serialized: {yaml}"
734 );
735 }
736
737 #[test]
738 fn role_entry_parses_route_policy() {
739 let yaml = r#"
740schema_version: 1
741models:
742 haiku:
743 provider: anthropic
744 model: claude-haiku-4-5
745 opus:
746 provider: anthropic
747 model: claude-opus-4-7
748roles:
749 dev:
750 primary: opus
751 route_policy: !force_frontier
752 model_id: opus
753 reflector:
754 primary: haiku
755 route_policy: prefer_local
756 curator:
757 primary: haiku
758 route_policy: force_local
759 chat:
760 primary: haiku
761"#;
762 let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
763 assert_eq!(
764 r.roles["dev"].route_policy,
765 Some(RoutePolicy::ForceFrontier {
766 model_id: "opus".into()
767 })
768 );
769 assert_eq!(
770 r.roles["reflector"].route_policy,
771 Some(RoutePolicy::PreferLocal)
772 );
773 assert_eq!(
774 r.roles["curator"].route_policy,
775 Some(RoutePolicy::ForceLocal)
776 );
777 assert_eq!(r.roles["chat"].route_policy, None);
778 }
779
780 #[test]
781 fn parses_split_cost_fields() {
782 let yaml = r#"
783schema_version: 1
784models:
785 opus:
786 provider: anthropic
787 model: claude-opus-4-8
788 input_cost_per_1k: 0.005
789 output_cost_per_1k: 0.025
790 context_window: 200000
791"#;
792 let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
793 let e = r.models.get("opus").unwrap();
794 assert_eq!(e.input_cost_per_1k, Some(0.005));
795 assert_eq!(e.output_cost_per_1k, Some(0.025));
796 assert_eq!(e.context_window, Some(200_000));
797 }
798
799 #[test]
800 fn default_model_entry_is_empty() {
801 let e = ModelEntry::default();
802 assert!(e.provider.is_empty());
803 assert_eq!(e.input_cost_per_1k, None);
804 assert_eq!(e.output_cost_per_1k, None);
805 assert_eq!(e.context_window, None);
806 }
807
808 #[test]
809 fn effective_costs_fallback_matrix() {
810 let mut e = ModelEntry {
812 cost_per_1k_tokens: Some(0.01),
813 ..Default::default()
814 };
815 assert_eq!(e.effective_costs(), (Some(0.01), Some(0.01)));
816
817 e = ModelEntry {
819 input_cost_per_1k: Some(0.005),
820 output_cost_per_1k: Some(0.025),
821 ..Default::default()
822 };
823 assert_eq!(e.effective_costs(), (Some(0.005), Some(0.025)));
824
825 e = ModelEntry {
827 cost_per_1k_tokens: Some(0.01),
828 input_cost_per_1k: Some(0.005),
829 output_cost_per_1k: Some(0.025),
830 ..Default::default()
831 };
832 assert_eq!(e.effective_costs(), (Some(0.005), Some(0.025)));
833
834 e = ModelEntry::default();
836 assert_eq!(e.effective_costs(), (None, None));
837 }
838}
839
840#[cfg(test)]
841mod io_tests {
842 use super::*;
843 use tempfile::tempdir;
844
845 #[test]
846 fn load_returns_empty_when_file_missing() {
847 let dir = tempdir().unwrap();
848 let r = ModelRegistry::load_from(&dir.path().join("nope.yaml")).unwrap();
849 assert_eq!(r.models.len(), 0);
850 assert_eq!(r.schema_version, 1);
851 }
852
853 #[test]
854 fn save_then_load_round_trips() {
855 let dir = tempdir().unwrap();
856 let p = dir.path().join("models.yaml");
857 let mut r = ModelRegistry::default();
858 r.models.insert(
859 "x".into(),
860 ModelEntry {
861 provider: "ollama".into(),
862 model: "llama3.2:3b".into(),
863 base_url: None,
864 secret: None,
865 capabilities: vec![],
866 params: serde_json::Value::Null,
867 tier: None,
868 cost_per_1k_tokens: None,
869 input_cost_per_1k: None,
870 output_cost_per_1k: None,
871 context_window: None,
872 priced_at: None,
873 ..Default::default()
874 },
875 );
876 r.save_to(&p).unwrap();
877 let r2 = ModelRegistry::load_from(&p).unwrap();
878 assert_eq!(r, r2);
879 }
880
881 #[test]
882 fn save_uses_atomic_rename() {
883 let dir = tempdir().unwrap();
884 let p = dir.path().join("models.yaml");
885 ModelRegistry::default().save_to(&p).unwrap();
886 let temp = dir.path().join("models.yaml.tmp");
887 assert!(!temp.exists(), "atomic temp left behind");
888 }
889}
890
891#[cfg(test)]
892mod switch_tests {
893 use super::*;
894 use crate::agent::AgentProfile;
895 use crate::config::{ModelSwitchConfig, RoutingConfig};
896
897 fn profile(model_ref: Option<&str>, chain: &[&str]) -> AgentProfile {
898 let mut p = AgentProfile::default_for_tests();
899 p.model_ref = model_ref.map(|s| s.to_string());
900 p.fallback_chain = chain.iter().map(|s| s.to_string()).collect();
901 p
902 }
903
904 #[test]
905 fn per_agent_primary_and_chain_win_over_global() {
906 let cfg = ModelSwitchConfig {
907 default: Some("global_default".into()),
908 fallback_chain: vec!["g1".into(), "g2".into()],
909 ..Default::default()
910 };
911 let p = profile(Some("agent_primary"), &["agent_primary", "agent_fb"]);
912 assert_eq!(
914 resolve_model_refs(&p, &cfg, None),
915 vec!["agent_primary", "agent_fb"]
916 );
917 }
918
919 #[test]
920 fn falls_back_to_global_default_and_chain() {
921 let cfg = ModelSwitchConfig {
922 default: Some("global_default".into()),
923 fallback_chain: vec!["g1".into(), "global_default".into()],
924 ..Default::default()
925 };
926 let p = profile(None, &[]); assert_eq!(
929 resolve_model_refs(&p, &cfg, None),
930 vec!["global_default", "g1"]
931 );
932 }
933
934 #[test]
935 fn routed_primary_overrides_model_ref() {
936 let cfg = ModelSwitchConfig {
937 fallback_chain: vec!["g1".into()],
938 ..Default::default()
939 };
940 let p = profile(Some("agent_primary"), &[]);
941 assert_eq!(
942 resolve_model_refs(&p, &cfg, Some("frontier".into())),
943 vec!["frontier", "g1"]
944 );
945 }
946
947 #[test]
948 fn no_config_no_agent_yields_empty() {
949 let cfg = ModelSwitchConfig::default();
951 assert!(resolve_model_refs(&profile(None, &[]), &cfg, None).is_empty());
952 }
953
954 #[test]
955 fn difficulty_picks_frontier_over_threshold() {
956 let r = RoutingConfig {
957 enabled: true,
958 cheap: Some("cheap".into()),
959 frontier: Some("frontier".into()),
960 threshold_input_tokens: Some(1000),
961 };
962 assert_eq!(choose_by_difficulty(1500, &r), Some("frontier".into()));
963 assert_eq!(choose_by_difficulty(500, &r), Some("cheap".into()));
964 let bad = RoutingConfig {
966 enabled: true,
967 cheap: Some("c".into()),
968 frontier: None,
969 threshold_input_tokens: None,
970 };
971 assert_eq!(choose_by_difficulty(9999, &bad), None);
972 }
973
974 #[test]
975 fn pick_cheap_model_lowest_cost_chat_excluding_primary() {
976 let mut reg = ModelRegistry::default();
977 let mk = |cost: f64, caps: &[&str]| ModelEntry {
978 provider: "x".into(),
979 model: "m".into(),
980 capabilities: caps.iter().map(|s| s.to_string()).collect(),
981 cost_per_1k_tokens: Some(cost),
982 ..Default::default()
983 };
984 reg.models.insert("frontier".into(), mk(0.01, &["chat"]));
985 reg.models.insert("cheap".into(), mk(0.0001, &["chat"]));
986 reg.models
987 .insert("embed".into(), mk(0.00001, &["embedding"])); assert_eq!(
990 pick_cheap_model(®, Some("cheap"), &[]),
991 Some("frontier".into())
992 ); assert_eq!(pick_cheap_model(®, None, &[]), Some("cheap".into()));
994 let mut empty = ModelRegistry::default();
996 empty.models.insert("e".into(), mk(0.0, &["embedding"]));
997 assert_eq!(pick_cheap_model(&empty, None, &[]), None);
998 }
999
1000 #[test]
1001 fn satisfies_is_permissive_at_baseline_and_fail_closed_above_it() {
1002 let mk = |caps: &[&str]| ModelEntry {
1003 provider: "x".into(),
1004 model: "m".into(),
1005 capabilities: caps.iter().map(|s| s.to_string()).collect(),
1006 ..Default::default()
1007 };
1008 assert!(satisfies(&mk(&[]), &[]));
1010 assert!(satisfies(&mk(&["chat"]), &[]));
1011 assert!(!satisfies(&mk(&["embedding"]), &[]));
1012 assert!(!satisfies(&mk(&[]), &[Requirement::Vision]));
1014 assert!(!satisfies(&mk(&["chat"]), &[Requirement::Vision]));
1015 assert!(satisfies(&mk(&["chat", "vision"]), &[Requirement::Vision]));
1016 assert!(!satisfies(&mk(&["chat", "vision"]), &[Requirement::Tools]));
1017 assert!(satisfies(
1018 &mk(&["chat", "vision", "tools"]),
1019 &[Requirement::Vision, Requirement::Tools]
1020 ));
1021 }
1022
1023 #[test]
1030 fn undeclared_capabilities_pass_tools_but_never_vision() {
1031 let mk = |caps: &[&str]| ModelEntry {
1032 provider: "x".into(),
1033 model: "m".into(),
1034 capabilities: caps.iter().map(|s| s.to_string()).collect(),
1035 ..Default::default()
1036 };
1037 assert!(satisfies(&mk(&[]), &[Requirement::Tools]));
1039 assert!(!satisfies(&mk(&[]), &[Requirement::Vision]));
1040 assert!(!satisfies(
1041 &mk(&[]),
1042 &[Requirement::Vision, Requirement::Tools]
1043 ));
1044 assert!(!satisfies(&mk(&["chat"]), &[Requirement::Tools]));
1046 assert!(satisfies(&mk(&["chat", "tools"]), &[Requirement::Tools]));
1047 }
1048
1049 #[test]
1052 fn pick_cheap_model_declines_when_no_entry_declares_the_requirement() {
1053 let mk = |cost: f64, caps: &[&str]| ModelEntry {
1054 provider: "x".into(),
1055 model: "m".into(),
1056 capabilities: caps.iter().map(|s| s.to_string()).collect(),
1057 cost_per_1k_tokens: Some(cost),
1058 ..Default::default()
1059 };
1060 let mut reg = ModelRegistry::default();
1061 reg.models
1062 .insert("cheap_text".into(), mk(0.0001, &["chat"]));
1063 reg.models.insert("legacy".into(), mk(0.0002, &[]));
1064 reg.models
1065 .insert("frontier".into(), mk(0.01, &["chat", "vision"]));
1066 assert_eq!(pick_cheap_model(®, None, &[]), Some("cheap_text".into()));
1068 assert_eq!(
1070 pick_cheap_model(®, None, &[Requirement::Vision]),
1071 Some("frontier".into())
1072 );
1073 let mut blind = ModelRegistry::default();
1075 blind
1076 .models
1077 .insert("cheap_text".into(), mk(0.0001, &["chat"]));
1078 blind.models.insert("legacy".into(), mk(0.0002, &[]));
1079 assert_eq!(pick_cheap_model(&blind, None, &[Requirement::Vision]), None);
1080 }
1081
1082 #[test]
1086 fn priced_at_stamps_only_priced_entries() {
1087 let now = chrono::Utc::now();
1088
1089 let mut unpriced = ModelEntry {
1090 provider: "openai".into(),
1091 model: "local-thing".into(),
1092 ..Default::default()
1093 };
1094 unpriced.stamp_priced_at(now);
1095 assert_eq!(unpriced.priced_at, None);
1096 assert_eq!(unpriced.price_age(now), None);
1097
1098 let mut priced = ModelEntry {
1099 output_cost_per_1k: Some(0.025),
1100 ..unpriced.clone()
1101 };
1102 priced.stamp_priced_at(now);
1103 assert_eq!(priced.priced_at, Some(now));
1104
1105 let mut legacy = ModelEntry {
1107 cost_per_1k_tokens: Some(0.01),
1108 ..unpriced.clone()
1109 };
1110 legacy.stamp_priced_at(now);
1111 assert!(legacy.priced_at.is_some());
1112
1113 let earlier = now - chrono::TimeDelta::days(30);
1115 priced.stamp_priced_at(earlier);
1116 assert_eq!(priced.priced_at, Some(now));
1117 }
1118
1119 #[test]
1122 fn registry_without_priced_at_still_loads_and_reports_unknown_age() {
1123 let yaml = r#"
1124schema_version: 1
1125models:
1126 opus:
1127 provider: anthropic
1128 model: claude-opus-5
1129 input_cost_per_1k: 0.005
1130 output_cost_per_1k: 0.025
1131"#;
1132 let reg: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
1133 let e = ®.models["opus"];
1134 assert_eq!(e.priced_at, None);
1135 assert_eq!(e.price_age(chrono::Utc::now()), None);
1136 let out = serde_yaml_ng::to_string(®).unwrap();
1138 assert!(!out.contains("priced_at"), "{out}");
1139 }
1140
1141 #[test]
1144 fn pick_cheap_model_sees_split_cost_entries() {
1145 let mut reg = ModelRegistry::default();
1146 let split = |input: f64, output: f64| ModelEntry {
1147 provider: "x".into(),
1148 model: "m".into(),
1149 capabilities: vec!["chat".into()],
1150 input_cost_per_1k: Some(input),
1151 output_cost_per_1k: Some(output),
1152 ..Default::default()
1153 };
1154 reg.models.insert("dear".into(), split(0.005, 0.025));
1155 reg.models.insert("cheap".into(), split(0.0001, 0.0004));
1156 assert_eq!(pick_cheap_model(®, None, &[]), Some("cheap".into()));
1157
1158 let mut input_only = ModelRegistry::default();
1160 input_only.models.insert(
1161 "in".into(),
1162 ModelEntry {
1163 provider: "x".into(),
1164 model: "m".into(),
1165 capabilities: vec!["chat".into()],
1166 input_cost_per_1k: Some(0.002),
1167 ..Default::default()
1168 },
1169 );
1170 assert_eq!(pick_cheap_model(&input_only, None, &[]), Some("in".into()));
1171 }
1172
1173 #[test]
1174 fn subscription_metadata_round_trips_without_a_secret() {
1175 let yaml = r#"schema_version: 1
1176models:
1177 chatgpt_sol:
1178 provider: codex
1179 model: gpt-5.6-sol
1180 base_url: http://127.0.0.1:8088/codex/v1
1181 tier: frontier
1182 billing: subscription
1183 catalog_verified: true
1184"#;
1185 let reg: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
1186 let entry = ®.models["chatgpt_sol"];
1187 assert_eq!(entry.billing, Some(BillingMode::Subscription));
1188 assert_eq!(entry.catalog_verified, Some(true));
1189 assert!(entry.secret.is_none());
1190 let out = serde_yaml_ng::to_string(®).unwrap();
1191 assert!(out.contains("billing: subscription"), "{out}");
1192 assert!(out.contains("catalog_verified: true"), "{out}");
1193 }
1194
1195 #[test]
1202 fn provider_inference_and_route_tier_cover_existing_aliases() {
1203 for provider in [
1204 "ollama",
1205 "mlx",
1206 "llamacpp",
1207 "llama_cpp",
1208 "localai",
1209 "lmstudio",
1210 "local",
1211 ] {
1212 assert_eq!(
1213 inferred_billing_for_provider(provider),
1214 BillingMode::Local,
1215 "{provider}"
1216 );
1217 let entry = ModelEntry {
1218 provider: provider.into(),
1219 ..Default::default()
1220 };
1221 assert_eq!(entry.effective_route_tier(), RouteTier::Local, "{provider}");
1222 }
1223 for provider in ["claude", "codex"] {
1224 assert_eq!(
1225 inferred_billing_for_provider(provider),
1226 BillingMode::Subscription
1227 );
1228 }
1229 assert_eq!(
1230 inferred_billing_for_provider("unknown"),
1231 BillingMode::UsageBilled
1232 );
1233 }
1234
1235 #[test]
1236 fn projected_cost_uses_four_to_one_workload_and_legacy_rates() {
1237 let split = ModelEntry {
1238 input_cost_per_1k: Some(1.0),
1239 output_cost_per_1k: Some(10.0),
1240 ..Default::default()
1241 };
1242 assert_eq!(split.projected_cost(4_000, 1_000), Some(14.0));
1243 let legacy = ModelEntry {
1244 cost_per_1k_tokens: Some(2.0),
1245 ..Default::default()
1246 };
1247 assert_eq!(legacy.projected_cost(4_000, 1_000), Some(10.0));
1248 assert_eq!(ModelEntry::default().projected_cost(4_000, 1_000), None);
1249 }
1250
1251 #[test]
1252 fn pick_cheap_model_uses_projected_four_to_one_cost() {
1253 let mut reg = ModelRegistry::default();
1254 reg.models.insert(
1255 "cheap_input".into(),
1256 ModelEntry {
1257 provider: "openai".into(),
1258 model: "a".into(),
1259 input_cost_per_1k: Some(0.1),
1260 output_cost_per_1k: Some(2.0),
1261 ..Default::default()
1262 },
1263 );
1264 reg.models.insert(
1265 "cheap_output".into(),
1266 ModelEntry {
1267 provider: "openai".into(),
1268 model: "b".into(),
1269 input_cost_per_1k: Some(1.0),
1270 output_cost_per_1k: Some(0.1),
1271 ..Default::default()
1272 },
1273 );
1274 assert_eq!(
1275 pick_cheap_model(®, None, &[]),
1276 Some("cheap_input".into())
1277 );
1278 }
1279
1280 #[test]
1281 fn billing_is_inferred_from_the_provider_when_not_declared() {
1282 let mut e = ModelEntry {
1283 provider: "ollama".into(),
1284 model: "llama3.2:3b".into(),
1285 ..Default::default()
1286 };
1287 assert_eq!(e.billing_or_inferred(), BillingMode::Local);
1288 e.provider = "codex".into();
1289 assert_eq!(e.billing_or_inferred(), BillingMode::Subscription);
1290 e.provider = "claude".into();
1291 assert_eq!(e.billing_or_inferred(), BillingMode::Subscription);
1292 e.provider = "openai".into();
1293 assert_eq!(
1294 e.billing_or_inferred(),
1295 BillingMode::UsageBilled,
1296 "unknown is metered"
1297 );
1298 e.provider = "anthropic".into();
1299 assert_eq!(e.billing_or_inferred(), BillingMode::UsageBilled);
1300 e.billing = Some(BillingMode::Local);
1303 assert_eq!(e.billing_or_inferred(), BillingMode::Local);
1304 }
1305
1306 #[test]
1307 fn entry_without_billing_metadata_stays_unknown() {
1308 let yaml = r#"schema_version: 1
1309models:
1310 gpt:
1311 provider: openai
1312 model: gpt-4o
1313 secret: env:OPENAI_API_KEY
1314"#;
1315 let reg: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
1316 let entry = ®.models["gpt"];
1317 assert_eq!(entry.billing, None);
1318 assert_eq!(entry.catalog_verified, None);
1319 let out = serde_yaml_ng::to_string(®).unwrap();
1320 assert!(!out.contains("billing"), "{out}");
1321 assert!(!out.contains("catalog_verified"), "{out}");
1322 for (raw, mode) in [
1323 ("subscription", BillingMode::Subscription),
1324 ("usage_billed", BillingMode::UsageBilled),
1325 ("local", BillingMode::Local),
1326 ] {
1327 let m: BillingMode = serde_yaml_ng::from_str(raw).unwrap();
1328 assert_eq!(m, mode);
1329 }
1330 }
1331}