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/// Registry capability strings. The baseline (`chat`) is legacy-permissive —
294/// an entry with no `capabilities` at all predates the field and is assumed
295/// chat-capable. Everything above the baseline is fail-closed.
296pub const CAP_CHAT: &str = "chat";
297pub const CAP_TOOLS: &str = "tools";
298pub const CAP_VISION: &str = "vision";
299
300/// A capability the request needs from whatever model serves it. Derived from
301/// the request itself (an image in the messages, a tool list) and never from
302/// config: a router may only substitute a model that can do the job.
303#[derive(Debug, Clone, Copy, PartialEq, Eq)]
304pub enum Requirement {
305    /// The request carries an image; the model has to be able to see it.
306    Vision,
307    /// The request declares tools; the model has to be able to call them.
308    Tools,
309}
310
311impl Requirement {
312    /// The registry capability an entry must declare to satisfy this.
313    pub fn capability(self) -> &'static str {
314        match self {
315            Requirement::Vision => CAP_VISION,
316            Requirement::Tools => CAP_TOOLS,
317        }
318    }
319
320    /// Does an entry that declares NO capabilities at all satisfy this?
321    ///
322    /// The two requirements differ in how they fail, and the answer follows
323    /// the failure mode rather than a blanket rule:
324    ///
325    /// - `Vision`: **no**. A model that cannot see answers an image request
326    ///   with confident nonsense — silent, and unrecoverable for that turn.
327    ///   That is the failure this gate exists to prevent, so silence about
328    ///   vision is treated as absence of it.
329    /// - `Tools`: **yes**. A model that cannot call tools fails loudly (the
330    ///   provider rejects the request) and the existing retry/advance path
331    ///   already handles it. Treating undeclared as incapable would drop every
332    ///   entry written before `capabilities` existed — in practice most of a
333    ///   real registry — out of the fallback chain of every tool-carrying turn,
334    ///   which is a large regression bought for very little.
335    ///
336    /// An entry that DOES declare capabilities is taken at its word either
337    /// way: if it enumerated what it can do and left `tools` out, that is a
338    /// statement, not silence.
339    fn permitted_when_undeclared(self) -> bool {
340        match self {
341            Requirement::Vision => false,
342            Requirement::Tools => true,
343        }
344    }
345}
346
347/// Can this entry serve a request needing `reqs`?
348///
349/// No registry write path emits `vision` today, so a `Vision` requirement
350/// disqualifies every current entry — auto-substitution goes inert for image
351/// requests rather than answering them blind. The same code makes a finer
352/// distinction the day entries start declaring it; there is no second version
353/// of this function to write later.
354pub 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
368/// Pick the cheapest registry entry that can serve a request needing `reqs`,
369/// excluding `exclude` (the agent's own primary). None when no qualifying
370/// entry exists → caller keeps normal candidates (fail-expensive).
371pub 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            // Not the deprecated field directly: `mur model add --output-cost`
382            // deliberately leaves it unset, so reading it drops every entry
383            // added with the current flags instead of ranking it.
384            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        // Recorded vendor wins — this is what new entries carry.
397        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        // Legacy entry with no vendor: the endpoint host still identifies it,
406        // which is how registries written before the field keep working.
407        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        // Nothing to infer: provider is all there is.
415        let bare = ModelEntry {
416            provider: "anthropic".into(),
417            ..Default::default()
418        };
419        assert_eq!(bare.vendor_candidates(), vec!["anthropic"]);
420
421        // No duplicate when host and provider agree.
422        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(&reg).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        // Missing tier is None.
633        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        // legacy only → both fall back to the blended rate
733        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        // split only → split wins, legacy ignored
740        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        // both → split wins
748        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        // none → none
757        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        // per-agent model_ref is primary; per-agent chain used; primary de-duped.
835        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, &[]); // no per-agent model_ref or chain
849        // primary = global default; global chain used; primary de-duped out.
850        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        // Nothing configured → empty vec (caller falls back to inline model).
872        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        // Misconfigured (missing frontier) → None (fall through).
887        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"])); // not chat → skip
910        // cheapest chat-capable, excluding the agent's own primary:
911        assert_eq!(
912            pick_cheap_model(&reg, Some("cheap"), &[]),
913            Some("frontier".into())
914        ); // cheap excluded
915        assert_eq!(pick_cheap_model(&reg, None, &[]), Some("cheap".into()));
916        // no chat entries → None (Smart inert)
917        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        // Baseline: an entry written before the field existed is still chat.
931        assert!(satisfies(&mk(&[]), &[]));
932        assert!(satisfies(&mk(&["chat"]), &[]));
933        assert!(!satisfies(&mk(&["embedding"]), &[]));
934        // Above baseline: unstated is not permission.
935        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    /// Tools and Vision disagree about silence on purpose. A tool-incapable
946    /// model fails loudly and the chain advances; a blind one answers with
947    /// confident nonsense. So an entry that declares nothing keeps its place in
948    /// the chain for a tool turn — otherwise every pre-`capabilities` entry
949    /// (most of a real registry) would drop out of every tool-carrying request
950    /// — while the same silence disqualifies it for an image.
951    #[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        // Silence: permitted for tools, never for vision.
960        assert!(satisfies(&mk(&[]), &[Requirement::Tools]));
961        assert!(!satisfies(&mk(&[]), &[Requirement::Vision]));
962        assert!(!satisfies(
963            &mk(&[]),
964            &[Requirement::Vision, Requirement::Tools]
965        ));
966        // A declaration is taken at its word in both directions.
967        assert!(!satisfies(&mk(&["chat"]), &[Requirement::Tools]));
968        assert!(satisfies(&mk(&["chat", "tools"]), &[Requirement::Tools]));
969    }
970
971    /// The incident, as a regression test: an image request against a registry
972    /// where nothing declares vision must find no cheap candidate at all.
973    #[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        // No requirement -> cheapest wins (today's behaviour, unchanged).
989        assert_eq!(pick_cheap_model(&reg, None, &[]), Some("cheap_text".into()));
990        // Vision required -> only the declaring entry qualifies, cost be damned.
991        assert_eq!(
992            pick_cheap_model(&reg, None, &[Requirement::Vision]),
993            Some("frontier".into())
994        );
995        // Nothing declares vision -> None, so Smart goes inert.
996        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    /// A price with no date is a guess wearing the costume of a fact. But a
1005    /// date on an entry that carries no price would be the same lie in the
1006    /// other direction, so the stamp is conditional on there being a rate.
1007    #[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        // A legacy single-rate entry counts as priced.
1028        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        // Re-stamping never moves the date backwards.
1036        let earlier = now - chrono::TimeDelta::days(30);
1037        priced.stamp_priced_at(earlier);
1038        assert_eq!(priced.priced_at, Some(now));
1039    }
1040
1041    /// Entries written before this field existed must keep loading, and must
1042    /// report an unknown age rather than inheriting today's date.
1043    #[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 = &reg.models["opus"];
1056        assert_eq!(e.priced_at, None);
1057        assert_eq!(e.price_age(chrono::Utc::now()), None);
1058        // Round-trips without inventing the field.
1059        let out = serde_yaml_ng::to_string(&reg).unwrap();
1060        assert!(!out.contains("priced_at"), "{out}");
1061    }
1062
1063    /// `mur model add --input-cost/--output-cost` leaves `cost_per_1k_tokens`
1064    /// unset, so an entry priced the current way must still be rankable.
1065    #[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(&reg, None, &[]), Some("cheap".into()));
1079
1080        // Input-only entries are priced too, rather than silently skipped.
1081        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}