Skip to main content

mur_common/
model.rs

1//! Named model registry shared by all agents.
2//!
3//! On disk: `~/.mur/models.yaml`. Schema:
4//!
5//! ```yaml
6//! schema_version: 1
7//! models:
8//!   anthropic_opus_4_7:
9//!     provider: anthropic
10//!     model: claude-opus-4-7
11//!     secret: env:ANTHROPIC_API_KEY
12//!     capabilities: [chat, tools]
13//! ```
14
15use 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    /// Who makes this model — the models.dev catalog vendor (`deepseek`,
26    /// `groq`, `mistral`, …).
27    ///
28    /// Distinct from `provider`, which is the wire protocol MUR dials: a
29    /// DeepSeek entry is `provider: openai` + `vendor: deepseek`, because the
30    /// runtime reaches it over the OpenAI protocol while the catalog files it
31    /// under DeepSeek. Only recorded when the two differ — for Anthropic,
32    /// OpenAI and Ollama the protocol already names the vendor.
33    ///
34    /// `None` on entries written before this field existed; readers should go
35    /// through [`ModelEntry::vendor_candidates`] rather than reading it raw.
36    #[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    /// Routing tier: cheap/local vs frontier/expensive.
49    /// When absent, the router infers based on provider.
50    #[serde(default, skip_serializing_if = "Option::is_none")]
51    pub tier: Option<RouteTier>,
52    /// Estimated USD cost per 1000 output tokens.
53    /// Used for ledger cost estimates.
54    #[serde(default, skip_serializing_if = "Option::is_none")]
55    pub cost_per_1k_tokens: Option<f64>,
56    /// Estimated USD cost per 1000 input tokens.
57    /// New field for split input/output cost tracking.
58    #[serde(default, skip_serializing_if = "Option::is_none")]
59    pub input_cost_per_1k: Option<f64>,
60    /// Estimated USD cost per 1000 output tokens.
61    /// New field for split input/output cost tracking.
62    #[serde(default, skip_serializing_if = "Option::is_none")]
63    pub output_cost_per_1k: Option<f64>,
64    /// Model context window size in tokens.
65    #[serde(default, skip_serializing_if = "Option::is_none")]
66    pub context_window: Option<u64>,
67    /// When the rates above were recorded.
68    ///
69    /// Vendors move prices; a rate written months ago is a guess wearing the
70    /// costume of a fact, and nothing else on this struct can tell the two
71    /// apart. `None` means unknown — entries predating this field, or hand-
72    /// written ones — which is honest rather than defaulting to "fresh".
73    #[serde(default, skip_serializing_if = "Option::is_none")]
74    pub priced_at: Option<chrono::DateTime<chrono::Utc>>,
75}
76
77/// Vendor label implied by an endpoint host: `https://api.deepseek.com/v1` →
78/// `deepseek`. Best-effort — a host that does not carry the vendor's name
79/// (Google's `generativelanguage.googleapis.com`) yields the wrong label,
80/// which is why `vendor` is recorded explicitly on new entries.
81fn 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    /// Resolve effective per-1k rates as `(input, output)`.
96    ///
97    /// The deprecated `cost_per_1k_tokens` is treated as the output rate and
98    /// also as the input fallback, so legacy single-rate entries keep working.
99    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    /// Catalog vendor names to try for this entry, most specific first.
106    ///
107    /// The recorded `vendor` wins. Failing that — legacy entries, or anything
108    /// written by hand — the host of `base_url` is tried
109    /// (`https://api.deepseek.com` → `deepseek`), then `provider`, which names
110    /// the vendor only when the vendor happens to have its own client.
111    ///
112    /// Every caller that asks an external catalog about an entry must go
113    /// through this. Asking with `provider` alone reports every
114    /// OpenAI-compatible third party as unknown.
115    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    /// Whether this entry carries any rate at all.
133    pub fn is_priced(&self) -> bool {
134        let (input, output) = self.effective_costs();
135        input.is_some() || output.is_some()
136    }
137
138    /// Stamp `priced_at` with `now`, but only if a rate is actually present —
139    /// a date on an unpriced entry would claim a freshness it does not have.
140    /// Never overwrites an existing stamp with an older one.
141    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    /// How long ago the rates were recorded, or `None` when unstamped.
148    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    /// Registry model ID (key in `models:`) to use as primary.
156    pub primary: String,
157    /// Fallback model ID if primary is unavailable.
158    #[serde(default, skip_serializing_if = "Option::is_none")]
159    pub fallback: Option<String>,
160    /// Optional daily cost cap in USD.
161    #[serde(default, skip_serializing_if = "Option::is_none")]
162    pub cost_budget_per_day_usd: Option<f64>,
163    /// If true, only use local models when handling sensitive data.
164    #[serde(default)]
165    pub privacy_local_only: bool,
166    /// Per-role routing policy override.
167    /// When absent, the router uses the default heuristic.
168    #[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        // Honor MUR_HOME (used by test harnesses and Windows CI, where
216        // `dirs::home_dir()` reads SHGetKnownFolderPath and ignores HOME).
217        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    /// Return the primary model ID for `role`, or the fallback if the primary
227    /// is not in the `models` map, or `None` if the role is not configured.
228    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        // primary not in registry — try fallback
234        if let Some(fb) = &entry.fallback
235            && self.models.contains_key(fb)
236        {
237            return Some(fb);
238        }
239        // role configured but no available model
240        None
241    }
242}
243
244use crate::agent::AgentProfile;
245use crate::config::{DEFAULT_ROUTING_THRESHOLD, ModelSwitchConfig, RoutingConfig};
246
247/// Build the ordered list of model_refs to try: `[primary, ...fallback]`.
248/// Priority per-agent → global. The primary is de-duplicated out of the chain
249/// (no point retrying the same ref back-to-back). Returns empty when nothing is
250/// configured, so the caller keeps today's single-inline-model behaviour.
251pub 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
276/// Opt-in difficulty heuristic: pick `frontier` when the estimated input token
277/// count exceeds the threshold, else `cheap`. `None` when misconfigured (caller
278/// falls through to model_ref/global default).
279pub 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
293/// Pick the cheapest chat-capable model_ref for Smart background routing,
294/// excluding `exclude` (the agent's own primary). Chat-capable = capabilities
295/// contains "chat" OR is empty (legacy entries assumed chat). None when no
296/// qualifying entry exists → caller keeps normal candidates (fail-expensive).
297pub 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            // Not the deprecated field directly: `mur model add --output-cost`
304            // deliberately leaves it unset, so reading it drops every entry
305            // added with the current flags instead of ranking it.
306            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        // Recorded vendor wins — this is what new entries carry.
319        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        // Legacy entry with no vendor: the endpoint host still identifies it,
328        // which is how registries written before the field keep working.
329        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        // Nothing to infer: provider is all there is.
337        let bare = ModelEntry {
338            provider: "anthropic".into(),
339            ..Default::default()
340        };
341        assert_eq!(bare.vendor_candidates(), vec!["anthropic"]);
342
343        // No duplicate when host and provider agree.
344        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(&reg).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        // Missing tier is None.
555        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        // legacy only → both fall back to the blended rate
655        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        // split only → split wins, legacy ignored
662        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        // both → split wins
670        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        // none → none
679        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        // per-agent model_ref is primary; per-agent chain used; primary de-duped.
757        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, &[]); // no per-agent model_ref or chain
771        // primary = global default; global chain used; primary de-duped out.
772        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        // Nothing configured → empty vec (caller falls back to inline model).
794        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        // Misconfigured (missing frontier) → None (fall through).
810        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"])); // not chat → skip
834        // cheapest chat-capable, excluding the agent's own primary:
835        assert_eq!(
836            pick_cheap_model(&reg, Some("cheap")),
837            Some("frontier".into())
838        ); // cheap excluded
839        assert_eq!(pick_cheap_model(&reg, None), Some("cheap".into()));
840        // no chat entries → None (Smart inert)
841        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    /// A price with no date is a guess wearing the costume of a fact. But a
847    /// date on an entry that carries no price would be the same lie in the
848    /// other direction, so the stamp is conditional on there being a rate.
849    #[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        // A legacy single-rate entry counts as priced.
870        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        // Re-stamping never moves the date backwards.
878        let earlier = now - chrono::TimeDelta::days(30);
879        priced.stamp_priced_at(earlier);
880        assert_eq!(priced.priced_at, Some(now));
881    }
882
883    /// Entries written before this field existed must keep loading, and must
884    /// report an unknown age rather than inheriting today's date.
885    #[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 = &reg.models["opus"];
898        assert_eq!(e.priced_at, None);
899        assert_eq!(e.price_age(chrono::Utc::now()), None);
900        // Round-trips without inventing the field.
901        let out = serde_yaml_ng::to_string(&reg).unwrap();
902        assert!(!out.contains("priced_at"), "{out}");
903    }
904
905    /// `mur model add --input-cost/--output-cost` leaves `cost_per_1k_tokens`
906    /// unset, so an entry priced the current way must still be rankable.
907    #[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(&reg, None), Some("cheap".into()));
921
922        // Input-only entries are priced too, rather than silently skipped.
923        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}