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/// Who pays when this model answers.
22///
23/// A ChatGPT-subscription model (`provider: codex`) and an OpenAI Platform
24/// model can share a model id and a wire format while landing on different
25/// bills, so the registry says which. `None` on entries written before this
26/// field existed — readers render that as unknown, never as free.
27#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
28#[serde(rename_all = "snake_case")]
29pub enum BillingMode {
30    /// Covered by a flat subscription (ChatGPT Plus/Pro via Codex).
31    Subscription,
32    /// Metered per token against an API key.
33    UsageBilled,
34    /// Runs on this machine; no bill.
35    Local,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
39pub struct ModelEntry {
40    #[serde(default)]
41    pub provider: String,
42    /// Who makes this model — the models.dev catalog vendor (`deepseek`,
43    /// `groq`, `mistral`, …).
44    ///
45    /// Distinct from `provider`, which is the wire protocol MUR dials: a
46    /// DeepSeek entry is `provider: openai` + `vendor: deepseek`, because the
47    /// runtime reaches it over the OpenAI protocol while the catalog files it
48    /// under DeepSeek. Only recorded when the two differ — for Anthropic,
49    /// OpenAI and Ollama the protocol already names the vendor.
50    ///
51    /// `None` on entries written before this field existed; readers should go
52    /// through [`ModelEntry::vendor_candidates`] rather than reading it raw.
53    #[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    /// Routing tier: cheap/local vs frontier/expensive.
66    /// When absent, the router infers based on provider.
67    #[serde(default, skip_serializing_if = "Option::is_none")]
68    pub tier: Option<RouteTier>,
69    /// Estimated USD cost per 1000 output tokens.
70    /// Used for ledger cost estimates.
71    #[serde(default, skip_serializing_if = "Option::is_none")]
72    pub cost_per_1k_tokens: Option<f64>,
73    /// Estimated USD cost per 1000 input tokens.
74    /// New field for split input/output cost tracking.
75    #[serde(default, skip_serializing_if = "Option::is_none")]
76    pub input_cost_per_1k: Option<f64>,
77    /// Estimated USD cost per 1000 output tokens.
78    /// New field for split input/output cost tracking.
79    #[serde(default, skip_serializing_if = "Option::is_none")]
80    pub output_cost_per_1k: Option<f64>,
81    /// Model context window size in tokens.
82    #[serde(default, skip_serializing_if = "Option::is_none")]
83    pub context_window: Option<u64>,
84    /// When the rates above were recorded.
85    ///
86    /// Vendors move prices; a rate written months ago is a guess wearing the
87    /// costume of a fact, and nothing else on this struct can tell the two
88    /// apart. `None` means unknown — entries predating this field, or hand-
89    /// written ones — which is honest rather than defaulting to "fresh".
90    #[serde(default, skip_serializing_if = "Option::is_none")]
91    pub priced_at: Option<chrono::DateTime<chrono::Utc>>,
92    /// See [`BillingMode`]. `None` = unknown.
93    #[serde(default, skip_serializing_if = "Option::is_none")]
94    pub billing: Option<BillingMode>,
95    /// Whether the model id came from the provider's live catalog
96    /// (`Some(true)`) or was typed by hand when discovery failed
97    /// (`Some(false)`). `None` on entries that predate the field.
98    #[serde(default, skip_serializing_if = "Option::is_none")]
99    pub catalog_verified: Option<bool>,
100}
101
102/// Vendor label implied by an endpoint host: `https://api.deepseek.com/v1` →
103/// `deepseek`. Best-effort — a host that does not carry the vendor's name
104/// (Google's `generativelanguage.googleapis.com`) yields the wrong label,
105/// which is why `vendor` is recorded explicitly on new entries.
106fn 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
119impl ModelEntry {
120    /// Resolve effective per-1k rates as `(input, output)`.
121    ///
122    /// The deprecated `cost_per_1k_tokens` is treated as the output rate and
123    /// also as the input fallback, so legacy single-rate entries keep working.
124    pub fn effective_costs(&self) -> (Option<f64>, Option<f64>) {
125        let output = self.output_cost_per_1k.or(self.cost_per_1k_tokens);
126        let input = self.input_cost_per_1k.or(self.cost_per_1k_tokens);
127        (input, output)
128    }
129
130    /// Catalog vendor names to try for this entry, most specific first.
131    ///
132    /// The recorded `vendor` wins. Failing that — legacy entries, or anything
133    /// written by hand — the host of `base_url` is tried
134    /// (`https://api.deepseek.com` → `deepseek`), then `provider`, which names
135    /// the vendor only when the vendor happens to have its own client.
136    ///
137    /// Every caller that asks an external catalog about an entry must go
138    /// through this. Asking with `provider` alone reports every
139    /// OpenAI-compatible third party as unknown.
140    pub fn vendor_candidates(&self) -> Vec<String> {
141        let mut out: Vec<String> = Vec::with_capacity(3);
142        let mut push = |v: &str| {
143            if !v.is_empty() && !out.iter().any(|e| e == v) {
144                out.push(v.to_string());
145            }
146        };
147        if let Some(v) = self.vendor.as_deref() {
148            push(v);
149        }
150        if let Some(label) = vendor_label_of_url(self.base_url.as_deref()) {
151            push(&label);
152        }
153        push(&self.provider);
154        out
155    }
156
157    /// Whether this entry carries any rate at all.
158    pub fn is_priced(&self) -> bool {
159        let (input, output) = self.effective_costs();
160        input.is_some() || output.is_some()
161    }
162
163    /// Stamp `priced_at` with `now`, but only if a rate is actually present —
164    /// a date on an unpriced entry would claim a freshness it does not have.
165    /// Never overwrites an existing stamp with an older one.
166    pub fn stamp_priced_at(&mut self, now: chrono::DateTime<chrono::Utc>) {
167        if self.is_priced() && self.priced_at.is_none_or(|prev| prev < now) {
168            self.priced_at = Some(now);
169        }
170    }
171
172    /// How long ago the rates were recorded, or `None` when unstamped.
173    pub fn price_age(&self, now: chrono::DateTime<chrono::Utc>) -> Option<chrono::TimeDelta> {
174        self.priced_at.map(|at| now - at)
175    }
176}
177
178#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
179pub struct RoleEntry {
180    /// Registry model ID (key in `models:`) to use as primary.
181    pub primary: String,
182    /// Fallback model ID if primary is unavailable.
183    #[serde(default, skip_serializing_if = "Option::is_none")]
184    pub fallback: Option<String>,
185    /// Optional daily cost cap in USD.
186    #[serde(default, skip_serializing_if = "Option::is_none")]
187    pub cost_budget_per_day_usd: Option<f64>,
188    /// If true, only use local models when handling sensitive data.
189    #[serde(default)]
190    pub privacy_local_only: bool,
191    /// Per-role routing policy override.
192    /// When absent, the router uses the default heuristic.
193    #[serde(default, skip_serializing_if = "Option::is_none")]
194    pub route_policy: Option<RoutePolicy>,
195}
196
197#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
198pub struct ModelRegistry {
199    pub schema_version: u32,
200    #[serde(default)]
201    pub models: BTreeMap<String, ModelEntry>,
202    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
203    pub roles: BTreeMap<String, RoleEntry>,
204}
205
206impl Default for ModelRegistry {
207    fn default() -> Self {
208        Self {
209            schema_version: 1,
210            models: BTreeMap::new(),
211            roles: BTreeMap::new(),
212        }
213    }
214}
215
216impl ModelRegistry {
217    pub fn load_from(path: &Path) -> anyhow::Result<Self> {
218        if !path.exists() {
219            return Ok(Self::default());
220        }
221        let body = std::fs::read_to_string(path)?;
222        if body.trim().is_empty() {
223            return Ok(Self::default());
224        }
225        Ok(serde_yaml_ng::from_str(&body)?)
226    }
227
228    pub fn save_to(&self, path: &Path) -> anyhow::Result<()> {
229        if let Some(parent) = path.parent() {
230            std::fs::create_dir_all(parent)?;
231        }
232        let body = serde_yaml_ng::to_string(self)?;
233        let tmp = path.with_extension("yaml.tmp");
234        std::fs::write(&tmp, body)?;
235        std::fs::rename(&tmp, path)?;
236        Ok(())
237    }
238
239    pub fn default_path() -> anyhow::Result<PathBuf> {
240        // Honor MUR_HOME (used by test harnesses and Windows CI, where
241        // `dirs::home_dir()` reads SHGetKnownFolderPath and ignores HOME).
242        if let Ok(p) = std::env::var("MUR_HOME")
243            && !p.is_empty()
244        {
245            return Ok(PathBuf::from(p).join("models.yaml"));
246        }
247        let home = dirs::home_dir().ok_or_else(|| anyhow::anyhow!("no home dir"))?;
248        Ok(home.join(".mur/models.yaml"))
249    }
250
251    /// Return the primary model ID for `role`, or the fallback if the primary
252    /// is not in the `models` map, or `None` if the role is not configured.
253    pub fn resolve_role(&self, role: &str) -> Option<&str> {
254        let entry = self.roles.get(role)?;
255        if self.models.contains_key(&entry.primary) {
256            return Some(&entry.primary);
257        }
258        // primary not in registry — try fallback
259        if let Some(fb) = &entry.fallback
260            && self.models.contains_key(fb)
261        {
262            return Some(fb);
263        }
264        // role configured but no available model
265        None
266    }
267}
268
269use crate::agent::AgentProfile;
270use crate::config::{DEFAULT_ROUTING_THRESHOLD, ModelSwitchConfig, RoutingConfig};
271
272/// Build the ordered list of model_refs to try: `[primary, ...fallback]`.
273/// Priority per-agent → global. The primary is de-duplicated out of the chain
274/// (no point retrying the same ref back-to-back). Returns empty when nothing is
275/// configured, so the caller keeps today's single-inline-model behaviour.
276pub fn resolve_model_refs(
277    profile: &AgentProfile,
278    cfg: &ModelSwitchConfig,
279    routed_primary: Option<String>,
280) -> Vec<String> {
281    let primary = routed_primary
282        .or_else(|| profile.model_ref.clone())
283        .or_else(|| cfg.default.clone());
284    let chain = if !profile.fallback_chain.is_empty() {
285        profile.fallback_chain.clone()
286    } else {
287        cfg.fallback_chain.clone()
288    };
289    let mut out: Vec<String> = Vec::new();
290    if let Some(p) = primary {
291        out.push(p);
292    }
293    for r in chain {
294        if !out.contains(&r) {
295            out.push(r);
296        }
297    }
298    out
299}
300
301/// Opt-in difficulty heuristic: pick `frontier` when the estimated input token
302/// count exceeds the threshold, else `cheap`. `None` when misconfigured (caller
303/// falls through to model_ref/global default).
304pub fn choose_by_difficulty(est_input_tokens: u32, r: &RoutingConfig) -> Option<String> {
305    let threshold = r
306        .threshold_input_tokens
307        .unwrap_or(DEFAULT_ROUTING_THRESHOLD);
308    match (r.cheap.as_ref(), r.frontier.as_ref()) {
309        (Some(cheap), Some(frontier)) => Some(if est_input_tokens > threshold {
310            frontier.clone()
311        } else {
312            cheap.clone()
313        }),
314        _ => None,
315    }
316}
317
318/// Registry capability strings. The baseline (`chat`) is legacy-permissive —
319/// an entry with no `capabilities` at all predates the field and is assumed
320/// chat-capable. Everything above the baseline is fail-closed.
321pub const CAP_CHAT: &str = "chat";
322pub const CAP_TOOLS: &str = "tools";
323pub const CAP_VISION: &str = "vision";
324
325/// A capability the request needs from whatever model serves it. Derived from
326/// the request itself (an image in the messages, a tool list) and never from
327/// config: a router may only substitute a model that can do the job.
328#[derive(Debug, Clone, Copy, PartialEq, Eq)]
329pub enum Requirement {
330    /// The request carries an image; the model has to be able to see it.
331    Vision,
332    /// The request declares tools; the model has to be able to call them.
333    Tools,
334}
335
336impl Requirement {
337    /// The registry capability an entry must declare to satisfy this.
338    pub fn capability(self) -> &'static str {
339        match self {
340            Requirement::Vision => CAP_VISION,
341            Requirement::Tools => CAP_TOOLS,
342        }
343    }
344
345    /// Does an entry that declares NO capabilities at all satisfy this?
346    ///
347    /// The two requirements differ in how they fail, and the answer follows
348    /// the failure mode rather than a blanket rule:
349    ///
350    /// - `Vision`: **no**. A model that cannot see answers an image request
351    ///   with confident nonsense — silent, and unrecoverable for that turn.
352    ///   That is the failure this gate exists to prevent, so silence about
353    ///   vision is treated as absence of it.
354    /// - `Tools`: **yes**. A model that cannot call tools fails loudly (the
355    ///   provider rejects the request) and the existing retry/advance path
356    ///   already handles it. Treating undeclared as incapable would drop every
357    ///   entry written before `capabilities` existed — in practice most of a
358    ///   real registry — out of the fallback chain of every tool-carrying turn,
359    ///   which is a large regression bought for very little.
360    ///
361    /// An entry that DOES declare capabilities is taken at its word either
362    /// way: if it enumerated what it can do and left `tools` out, that is a
363    /// statement, not silence.
364    fn permitted_when_undeclared(self) -> bool {
365        match self {
366            Requirement::Vision => false,
367            Requirement::Tools => true,
368        }
369    }
370}
371
372/// Can this entry serve a request needing `reqs`?
373///
374/// No registry write path emits `vision` today, so a `Vision` requirement
375/// disqualifies every current entry — auto-substitution goes inert for image
376/// requests rather than answering them blind. The same code makes a finer
377/// distinction the day entries start declaring it; there is no second version
378/// of this function to write later.
379pub fn satisfies(e: &ModelEntry, reqs: &[Requirement]) -> bool {
380    let chat_capable = e.capabilities.is_empty() || e.capabilities.iter().any(|c| c == CAP_CHAT);
381    if !chat_capable {
382        return false;
383    }
384    reqs.iter().all(|r| {
385        if e.capabilities.is_empty() {
386            r.permitted_when_undeclared()
387        } else {
388            e.capabilities.iter().any(|c| c == r.capability())
389        }
390    })
391}
392
393/// Pick the cheapest registry entry that can serve a request needing `reqs`,
394/// excluding `exclude` (the agent's own primary). None when no qualifying
395/// entry exists → caller keeps normal candidates (fail-expensive).
396pub fn pick_cheap_model(
397    reg: &ModelRegistry,
398    exclude: Option<&str>,
399    reqs: &[Requirement],
400) -> Option<String> {
401    reg.models
402        .iter()
403        .filter(|(k, _)| exclude != Some(k.as_str()))
404        .filter(|(_, e)| satisfies(e, reqs))
405        .filter_map(|(k, e)| {
406            // Not the deprecated field directly: `mur model add --output-cost`
407            // deliberately leaves it unset, so reading it drops every entry
408            // added with the current flags instead of ranking it.
409            let (input, output) = e.effective_costs();
410            output.or(input).map(|c| (c, k.clone()))
411        })
412        .min_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal))
413        .map(|(_, k)| k)
414}
415
416#[cfg(test)]
417mod tests {
418
419    #[test]
420    fn vendor_candidates_prefer_the_recorded_vendor_then_the_host_then_provider() {
421        // Recorded vendor wins — this is what new entries carry.
422        let e = ModelEntry {
423            provider: "openai".into(),
424            vendor: Some("deepseek".into()),
425            base_url: Some("https://api.deepseek.com/v1".into()),
426            ..Default::default()
427        };
428        assert_eq!(e.vendor_candidates(), vec!["deepseek", "openai"]);
429
430        // Legacy entry with no vendor: the endpoint host still identifies it,
431        // which is how registries written before the field keep working.
432        let legacy = ModelEntry {
433            provider: "openai".into(),
434            base_url: Some("https://api.deepseek.com/v1".into()),
435            ..Default::default()
436        };
437        assert_eq!(legacy.vendor_candidates(), vec!["deepseek", "openai"]);
438
439        // Nothing to infer: provider is all there is.
440        let bare = ModelEntry {
441            provider: "anthropic".into(),
442            ..Default::default()
443        };
444        assert_eq!(bare.vendor_candidates(), vec!["anthropic"]);
445
446        // No duplicate when host and provider agree.
447        let same = ModelEntry {
448            provider: "openai".into(),
449            base_url: Some("https://api.openai.com/v1".into()),
450            ..Default::default()
451        };
452        assert_eq!(same.vendor_candidates(), vec!["openai"]);
453    }
454
455    #[test]
456    fn vendor_is_omitted_from_yaml_when_absent_and_round_trips_when_set() {
457        let bare = ModelEntry {
458            provider: "anthropic".into(),
459            model: "claude-opus-5".into(),
460            ..Default::default()
461        };
462        let y = serde_yaml_ng::to_string(&bare).unwrap();
463        assert!(!y.contains("vendor"), "{y}");
464
465        let tagged = ModelEntry {
466            provider: "openai".into(),
467            vendor: Some("groq".into()),
468            model: "llama-3.3".into(),
469            ..Default::default()
470        };
471        let y = serde_yaml_ng::to_string(&tagged).unwrap();
472        let back: ModelEntry = serde_yaml_ng::from_str(&y).unwrap();
473        assert_eq!(back.vendor.as_deref(), Some("groq"));
474    }
475    use super::*;
476
477    #[test]
478    fn parses_full_registry() {
479        let yaml = r#"
480schema_version: 1
481models:
482  anthropic_opus_4_7:
483    provider: anthropic
484    model: claude-opus-4-7
485    secret: env:ANTHROPIC_API_KEY
486    capabilities: [chat, tools]
487  ollama_llama3:
488    provider: ollama
489    model: llama3.2:3b
490    base_url: http://127.0.0.1:11434
491"#;
492        let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
493        assert_eq!(r.schema_version, 1);
494        assert_eq!(r.models.len(), 2);
495        let opus = r.models.get("anthropic_opus_4_7").unwrap();
496        assert_eq!(opus.provider, "anthropic");
497        assert_eq!(
498            opus.secret,
499            Some(SecretRef::Env("ANTHROPIC_API_KEY".into()))
500        );
501        assert!(r.models["ollama_llama3"].secret.is_none());
502    }
503
504    #[test]
505    fn round_trip_preserves_shape() {
506        let mut r = ModelRegistry::default();
507        r.models.insert(
508            "foo".into(),
509            ModelEntry {
510                provider: "anthropic".into(),
511                model: "claude-opus-4-7".into(),
512                base_url: None,
513                secret: Some(SecretRef::Keychain {
514                    service: "mur".into(),
515                    account: "anthropic".into(),
516                }),
517                capabilities: vec!["chat".into()],
518                params: serde_json::Value::Null,
519                tier: None,
520                cost_per_1k_tokens: None,
521                input_cost_per_1k: None,
522                output_cost_per_1k: None,
523                context_window: None,
524                priced_at: None,
525                ..Default::default()
526            },
527        );
528        let s = serde_yaml_ng::to_string(&r).unwrap();
529        let parsed: ModelRegistry = serde_yaml_ng::from_str(&s).unwrap();
530        assert_eq!(r, parsed);
531    }
532
533    #[test]
534    fn rejects_unknown_secret_scheme() {
535        let yaml = r#"
536schema_version: 1
537models:
538  bad:
539    provider: x
540    model: y
541    secret: bogus:value
542"#;
543        let r: Result<ModelRegistry, _> = serde_yaml_ng::from_str(yaml);
544        assert!(r.is_err(), "should reject unknown scheme");
545    }
546
547    #[test]
548    fn test_registry_roundtrip_with_roles() {
549        let yaml = r#"
550schema_version: 1
551models:
552  haiku:
553    provider: anthropic
554    model: claude-haiku-4-5
555roles:
556  reflector:
557    primary: haiku
558    fallback: null
559    cost_budget_per_day_usd: 0.5
560"#;
561        let reg: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
562        assert_eq!(reg.roles["reflector"].primary, "haiku");
563        let back = serde_yaml_ng::to_string(&reg).unwrap();
564        let reg2: ModelRegistry = serde_yaml_ng::from_str(&back).unwrap();
565        assert_eq!(reg, reg2);
566    }
567
568    #[test]
569    fn test_resolve_role_primary() {
570        let mut reg = ModelRegistry::default();
571        reg.models.insert(
572            "haiku".into(),
573            ModelEntry {
574                provider: "anthropic".into(),
575                model: "claude-haiku-4-5".into(),
576                base_url: None,
577                secret: None,
578                capabilities: vec![],
579                params: serde_json::Value::Null,
580                tier: None,
581                cost_per_1k_tokens: None,
582                input_cost_per_1k: None,
583                output_cost_per_1k: None,
584                context_window: None,
585                priced_at: None,
586                ..Default::default()
587            },
588        );
589        reg.roles.insert(
590            "reflector".into(),
591            RoleEntry {
592                primary: "haiku".into(),
593                fallback: None,
594                ..Default::default()
595            },
596        );
597        assert_eq!(reg.resolve_role("reflector"), Some("haiku"));
598    }
599
600    #[test]
601    fn test_resolve_role_fallback() {
602        let mut reg = ModelRegistry::default();
603        reg.models.insert(
604            "haiku".into(),
605            ModelEntry {
606                provider: "anthropic".into(),
607                model: "claude-haiku-4-5".into(),
608                base_url: None,
609                secret: None,
610                capabilities: vec![],
611                params: serde_json::Value::Null,
612                tier: None,
613                cost_per_1k_tokens: None,
614                input_cost_per_1k: None,
615                output_cost_per_1k: None,
616                context_window: None,
617                priced_at: None,
618                ..Default::default()
619            },
620        );
621        reg.roles.insert(
622            "reflector".into(),
623            RoleEntry {
624                primary: "nonexistent".into(),
625                fallback: Some("haiku".into()),
626                ..Default::default()
627            },
628        );
629        assert_eq!(reg.resolve_role("reflector"), Some("haiku"));
630    }
631
632    #[test]
633    fn test_resolve_role_none() {
634        let reg = ModelRegistry::default();
635        assert_eq!(reg.resolve_role("reflector"), None);
636    }
637
638    #[test]
639    fn model_entry_parses_tier_field() {
640        let yaml = r#"
641schema_version: 1
642models:
643  haiku:
644    provider: anthropic
645    model: claude-haiku-4-5
646    tier: local
647  opus:
648    provider: anthropic
649    model: claude-opus-4-7
650    tier: frontier
651    cost_per_1k_tokens: 0.015
652"#;
653        let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
654        assert_eq!(r.models["haiku"].tier, Some(RouteTier::Local));
655        assert_eq!(r.models["opus"].tier, Some(RouteTier::Frontier));
656        assert_eq!(r.models["opus"].cost_per_1k_tokens, Some(0.015));
657        // Missing tier is None.
658        let mut r2 = ModelRegistry::default();
659        r2.models.insert(
660            "x".into(),
661            ModelEntry {
662                provider: "ollama".into(),
663                model: "llama3".into(),
664                base_url: None,
665                secret: None,
666                capabilities: vec![],
667                params: serde_json::Value::Null,
668                tier: None,
669                cost_per_1k_tokens: None,
670                input_cost_per_1k: None,
671                output_cost_per_1k: None,
672                context_window: None,
673                priced_at: None,
674                ..Default::default()
675            },
676        );
677        let yaml = serde_yaml_ng::to_string(&r2).unwrap();
678        assert!(
679            !yaml.contains("tier:"),
680            "absent tier should not be serialized: {yaml}"
681        );
682    }
683
684    #[test]
685    fn role_entry_parses_route_policy() {
686        let yaml = r#"
687schema_version: 1
688models:
689  haiku:
690    provider: anthropic
691    model: claude-haiku-4-5
692  opus:
693    provider: anthropic
694    model: claude-opus-4-7
695roles:
696  dev:
697    primary: opus
698    route_policy: !force_frontier
699      model_id: opus
700  reflector:
701    primary: haiku
702    route_policy: prefer_local
703  curator:
704    primary: haiku
705    route_policy: force_local
706  chat:
707    primary: haiku
708"#;
709        let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
710        assert_eq!(
711            r.roles["dev"].route_policy,
712            Some(RoutePolicy::ForceFrontier {
713                model_id: "opus".into()
714            })
715        );
716        assert_eq!(
717            r.roles["reflector"].route_policy,
718            Some(RoutePolicy::PreferLocal)
719        );
720        assert_eq!(
721            r.roles["curator"].route_policy,
722            Some(RoutePolicy::ForceLocal)
723        );
724        assert_eq!(r.roles["chat"].route_policy, None);
725    }
726
727    #[test]
728    fn parses_split_cost_fields() {
729        let yaml = r#"
730schema_version: 1
731models:
732  opus:
733    provider: anthropic
734    model: claude-opus-4-8
735    input_cost_per_1k: 0.005
736    output_cost_per_1k: 0.025
737    context_window: 200000
738"#;
739        let r: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
740        let e = r.models.get("opus").unwrap();
741        assert_eq!(e.input_cost_per_1k, Some(0.005));
742        assert_eq!(e.output_cost_per_1k, Some(0.025));
743        assert_eq!(e.context_window, Some(200_000));
744    }
745
746    #[test]
747    fn default_model_entry_is_empty() {
748        let e = ModelEntry::default();
749        assert!(e.provider.is_empty());
750        assert_eq!(e.input_cost_per_1k, None);
751        assert_eq!(e.output_cost_per_1k, None);
752        assert_eq!(e.context_window, None);
753    }
754
755    #[test]
756    fn effective_costs_fallback_matrix() {
757        // legacy only → both fall back to the blended rate
758        let mut e = ModelEntry {
759            cost_per_1k_tokens: Some(0.01),
760            ..Default::default()
761        };
762        assert_eq!(e.effective_costs(), (Some(0.01), Some(0.01)));
763
764        // split only → split wins, legacy ignored
765        e = ModelEntry {
766            input_cost_per_1k: Some(0.005),
767            output_cost_per_1k: Some(0.025),
768            ..Default::default()
769        };
770        assert_eq!(e.effective_costs(), (Some(0.005), Some(0.025)));
771
772        // both → split wins
773        e = ModelEntry {
774            cost_per_1k_tokens: Some(0.01),
775            input_cost_per_1k: Some(0.005),
776            output_cost_per_1k: Some(0.025),
777            ..Default::default()
778        };
779        assert_eq!(e.effective_costs(), (Some(0.005), Some(0.025)));
780
781        // none → none
782        e = ModelEntry::default();
783        assert_eq!(e.effective_costs(), (None, None));
784    }
785}
786
787#[cfg(test)]
788mod io_tests {
789    use super::*;
790    use tempfile::tempdir;
791
792    #[test]
793    fn load_returns_empty_when_file_missing() {
794        let dir = tempdir().unwrap();
795        let r = ModelRegistry::load_from(&dir.path().join("nope.yaml")).unwrap();
796        assert_eq!(r.models.len(), 0);
797        assert_eq!(r.schema_version, 1);
798    }
799
800    #[test]
801    fn save_then_load_round_trips() {
802        let dir = tempdir().unwrap();
803        let p = dir.path().join("models.yaml");
804        let mut r = ModelRegistry::default();
805        r.models.insert(
806            "x".into(),
807            ModelEntry {
808                provider: "ollama".into(),
809                model: "llama3.2:3b".into(),
810                base_url: None,
811                secret: None,
812                capabilities: vec![],
813                params: serde_json::Value::Null,
814                tier: None,
815                cost_per_1k_tokens: None,
816                input_cost_per_1k: None,
817                output_cost_per_1k: None,
818                context_window: None,
819                priced_at: None,
820                ..Default::default()
821            },
822        );
823        r.save_to(&p).unwrap();
824        let r2 = ModelRegistry::load_from(&p).unwrap();
825        assert_eq!(r, r2);
826    }
827
828    #[test]
829    fn save_uses_atomic_rename() {
830        let dir = tempdir().unwrap();
831        let p = dir.path().join("models.yaml");
832        ModelRegistry::default().save_to(&p).unwrap();
833        let temp = dir.path().join("models.yaml.tmp");
834        assert!(!temp.exists(), "atomic temp left behind");
835    }
836}
837
838#[cfg(test)]
839mod switch_tests {
840    use super::*;
841    use crate::agent::AgentProfile;
842    use crate::config::{ModelSwitchConfig, RoutingConfig};
843
844    fn profile(model_ref: Option<&str>, chain: &[&str]) -> AgentProfile {
845        let mut p = AgentProfile::default_for_tests();
846        p.model_ref = model_ref.map(|s| s.to_string());
847        p.fallback_chain = chain.iter().map(|s| s.to_string()).collect();
848        p
849    }
850
851    #[test]
852    fn per_agent_primary_and_chain_win_over_global() {
853        let cfg = ModelSwitchConfig {
854            default: Some("global_default".into()),
855            fallback_chain: vec!["g1".into(), "g2".into()],
856            ..Default::default()
857        };
858        let p = profile(Some("agent_primary"), &["agent_primary", "agent_fb"]);
859        // per-agent model_ref is primary; per-agent chain used; primary de-duped.
860        assert_eq!(
861            resolve_model_refs(&p, &cfg, None),
862            vec!["agent_primary", "agent_fb"]
863        );
864    }
865
866    #[test]
867    fn falls_back_to_global_default_and_chain() {
868        let cfg = ModelSwitchConfig {
869            default: Some("global_default".into()),
870            fallback_chain: vec!["g1".into(), "global_default".into()],
871            ..Default::default()
872        };
873        let p = profile(None, &[]); // no per-agent model_ref or chain
874        // primary = global default; global chain used; primary de-duped out.
875        assert_eq!(
876            resolve_model_refs(&p, &cfg, None),
877            vec!["global_default", "g1"]
878        );
879    }
880
881    #[test]
882    fn routed_primary_overrides_model_ref() {
883        let cfg = ModelSwitchConfig {
884            fallback_chain: vec!["g1".into()],
885            ..Default::default()
886        };
887        let p = profile(Some("agent_primary"), &[]);
888        assert_eq!(
889            resolve_model_refs(&p, &cfg, Some("frontier".into())),
890            vec!["frontier", "g1"]
891        );
892    }
893
894    #[test]
895    fn no_config_no_agent_yields_empty() {
896        // Nothing configured → empty vec (caller falls back to inline model).
897        let cfg = ModelSwitchConfig::default();
898        assert!(resolve_model_refs(&profile(None, &[]), &cfg, None).is_empty());
899    }
900
901    #[test]
902    fn difficulty_picks_frontier_over_threshold() {
903        let r = RoutingConfig {
904            enabled: true,
905            cheap: Some("cheap".into()),
906            frontier: Some("frontier".into()),
907            threshold_input_tokens: Some(1000),
908        };
909        assert_eq!(choose_by_difficulty(1500, &r), Some("frontier".into()));
910        assert_eq!(choose_by_difficulty(500, &r), Some("cheap".into()));
911        // Misconfigured (missing frontier) → None (fall through).
912        let bad = RoutingConfig {
913            enabled: true,
914            cheap: Some("c".into()),
915            frontier: None,
916            threshold_input_tokens: None,
917        };
918        assert_eq!(choose_by_difficulty(9999, &bad), None);
919    }
920
921    #[test]
922    fn pick_cheap_model_lowest_cost_chat_excluding_primary() {
923        let mut reg = ModelRegistry::default();
924        let mk = |cost: f64, caps: &[&str]| ModelEntry {
925            provider: "x".into(),
926            model: "m".into(),
927            capabilities: caps.iter().map(|s| s.to_string()).collect(),
928            cost_per_1k_tokens: Some(cost),
929            ..Default::default()
930        };
931        reg.models.insert("frontier".into(), mk(0.01, &["chat"]));
932        reg.models.insert("cheap".into(), mk(0.0001, &["chat"]));
933        reg.models
934            .insert("embed".into(), mk(0.00001, &["embedding"])); // not chat → skip
935        // cheapest chat-capable, excluding the agent's own primary:
936        assert_eq!(
937            pick_cheap_model(&reg, Some("cheap"), &[]),
938            Some("frontier".into())
939        ); // cheap excluded
940        assert_eq!(pick_cheap_model(&reg, None, &[]), Some("cheap".into()));
941        // no chat entries → None (Smart inert)
942        let mut empty = ModelRegistry::default();
943        empty.models.insert("e".into(), mk(0.0, &["embedding"]));
944        assert_eq!(pick_cheap_model(&empty, None, &[]), None);
945    }
946
947    #[test]
948    fn satisfies_is_permissive_at_baseline_and_fail_closed_above_it() {
949        let mk = |caps: &[&str]| ModelEntry {
950            provider: "x".into(),
951            model: "m".into(),
952            capabilities: caps.iter().map(|s| s.to_string()).collect(),
953            ..Default::default()
954        };
955        // Baseline: an entry written before the field existed is still chat.
956        assert!(satisfies(&mk(&[]), &[]));
957        assert!(satisfies(&mk(&["chat"]), &[]));
958        assert!(!satisfies(&mk(&["embedding"]), &[]));
959        // Above baseline: unstated is not permission.
960        assert!(!satisfies(&mk(&[]), &[Requirement::Vision]));
961        assert!(!satisfies(&mk(&["chat"]), &[Requirement::Vision]));
962        assert!(satisfies(&mk(&["chat", "vision"]), &[Requirement::Vision]));
963        assert!(!satisfies(&mk(&["chat", "vision"]), &[Requirement::Tools]));
964        assert!(satisfies(
965            &mk(&["chat", "vision", "tools"]),
966            &[Requirement::Vision, Requirement::Tools]
967        ));
968    }
969
970    /// Tools and Vision disagree about silence on purpose. A tool-incapable
971    /// model fails loudly and the chain advances; a blind one answers with
972    /// confident nonsense. So an entry that declares nothing keeps its place in
973    /// the chain for a tool turn — otherwise every pre-`capabilities` entry
974    /// (most of a real registry) would drop out of every tool-carrying request
975    /// — while the same silence disqualifies it for an image.
976    #[test]
977    fn undeclared_capabilities_pass_tools_but_never_vision() {
978        let mk = |caps: &[&str]| ModelEntry {
979            provider: "x".into(),
980            model: "m".into(),
981            capabilities: caps.iter().map(|s| s.to_string()).collect(),
982            ..Default::default()
983        };
984        // Silence: permitted for tools, never for vision.
985        assert!(satisfies(&mk(&[]), &[Requirement::Tools]));
986        assert!(!satisfies(&mk(&[]), &[Requirement::Vision]));
987        assert!(!satisfies(
988            &mk(&[]),
989            &[Requirement::Vision, Requirement::Tools]
990        ));
991        // A declaration is taken at its word in both directions.
992        assert!(!satisfies(&mk(&["chat"]), &[Requirement::Tools]));
993        assert!(satisfies(&mk(&["chat", "tools"]), &[Requirement::Tools]));
994    }
995
996    /// The incident, as a regression test: an image request against a registry
997    /// where nothing declares vision must find no cheap candidate at all.
998    #[test]
999    fn pick_cheap_model_declines_when_no_entry_declares_the_requirement() {
1000        let mk = |cost: f64, caps: &[&str]| ModelEntry {
1001            provider: "x".into(),
1002            model: "m".into(),
1003            capabilities: caps.iter().map(|s| s.to_string()).collect(),
1004            cost_per_1k_tokens: Some(cost),
1005            ..Default::default()
1006        };
1007        let mut reg = ModelRegistry::default();
1008        reg.models
1009            .insert("cheap_text".into(), mk(0.0001, &["chat"]));
1010        reg.models.insert("legacy".into(), mk(0.0002, &[]));
1011        reg.models
1012            .insert("frontier".into(), mk(0.01, &["chat", "vision"]));
1013        // No requirement -> cheapest wins (today's behaviour, unchanged).
1014        assert_eq!(pick_cheap_model(&reg, None, &[]), Some("cheap_text".into()));
1015        // Vision required -> only the declaring entry qualifies, cost be damned.
1016        assert_eq!(
1017            pick_cheap_model(&reg, None, &[Requirement::Vision]),
1018            Some("frontier".into())
1019        );
1020        // Nothing declares vision -> None, so Smart goes inert.
1021        let mut blind = ModelRegistry::default();
1022        blind
1023            .models
1024            .insert("cheap_text".into(), mk(0.0001, &["chat"]));
1025        blind.models.insert("legacy".into(), mk(0.0002, &[]));
1026        assert_eq!(pick_cheap_model(&blind, None, &[Requirement::Vision]), None);
1027    }
1028
1029    /// A price with no date is a guess wearing the costume of a fact. But a
1030    /// date on an entry that carries no price would be the same lie in the
1031    /// other direction, so the stamp is conditional on there being a rate.
1032    #[test]
1033    fn priced_at_stamps_only_priced_entries() {
1034        let now = chrono::Utc::now();
1035
1036        let mut unpriced = ModelEntry {
1037            provider: "openai".into(),
1038            model: "local-thing".into(),
1039            ..Default::default()
1040        };
1041        unpriced.stamp_priced_at(now);
1042        assert_eq!(unpriced.priced_at, None);
1043        assert_eq!(unpriced.price_age(now), None);
1044
1045        let mut priced = ModelEntry {
1046            output_cost_per_1k: Some(0.025),
1047            ..unpriced.clone()
1048        };
1049        priced.stamp_priced_at(now);
1050        assert_eq!(priced.priced_at, Some(now));
1051
1052        // A legacy single-rate entry counts as priced.
1053        let mut legacy = ModelEntry {
1054            cost_per_1k_tokens: Some(0.01),
1055            ..unpriced.clone()
1056        };
1057        legacy.stamp_priced_at(now);
1058        assert!(legacy.priced_at.is_some());
1059
1060        // Re-stamping never moves the date backwards.
1061        let earlier = now - chrono::TimeDelta::days(30);
1062        priced.stamp_priced_at(earlier);
1063        assert_eq!(priced.priced_at, Some(now));
1064    }
1065
1066    /// Entries written before this field existed must keep loading, and must
1067    /// report an unknown age rather than inheriting today's date.
1068    #[test]
1069    fn registry_without_priced_at_still_loads_and_reports_unknown_age() {
1070        let yaml = r#"
1071schema_version: 1
1072models:
1073  opus:
1074    provider: anthropic
1075    model: claude-opus-5
1076    input_cost_per_1k: 0.005
1077    output_cost_per_1k: 0.025
1078"#;
1079        let reg: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
1080        let e = &reg.models["opus"];
1081        assert_eq!(e.priced_at, None);
1082        assert_eq!(e.price_age(chrono::Utc::now()), None);
1083        // Round-trips without inventing the field.
1084        let out = serde_yaml_ng::to_string(&reg).unwrap();
1085        assert!(!out.contains("priced_at"), "{out}");
1086    }
1087
1088    /// `mur model add --input-cost/--output-cost` leaves `cost_per_1k_tokens`
1089    /// unset, so an entry priced the current way must still be rankable.
1090    #[test]
1091    fn pick_cheap_model_sees_split_cost_entries() {
1092        let mut reg = ModelRegistry::default();
1093        let split = |input: f64, output: f64| ModelEntry {
1094            provider: "x".into(),
1095            model: "m".into(),
1096            capabilities: vec!["chat".into()],
1097            input_cost_per_1k: Some(input),
1098            output_cost_per_1k: Some(output),
1099            ..Default::default()
1100        };
1101        reg.models.insert("dear".into(), split(0.005, 0.025));
1102        reg.models.insert("cheap".into(), split(0.0001, 0.0004));
1103        assert_eq!(pick_cheap_model(&reg, None, &[]), Some("cheap".into()));
1104
1105        // Input-only entries are priced too, rather than silently skipped.
1106        let mut input_only = ModelRegistry::default();
1107        input_only.models.insert(
1108            "in".into(),
1109            ModelEntry {
1110                provider: "x".into(),
1111                model: "m".into(),
1112                capabilities: vec!["chat".into()],
1113                input_cost_per_1k: Some(0.002),
1114                ..Default::default()
1115            },
1116        );
1117        assert_eq!(pick_cheap_model(&input_only, None, &[]), Some("in".into()));
1118    }
1119
1120    #[test]
1121    fn subscription_metadata_round_trips_without_a_secret() {
1122        let yaml = r#"schema_version: 1
1123models:
1124  chatgpt_sol:
1125    provider: codex
1126    model: gpt-5.6-sol
1127    base_url: http://127.0.0.1:8088/codex/v1
1128    tier: frontier
1129    billing: subscription
1130    catalog_verified: true
1131"#;
1132        let reg: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
1133        let entry = &reg.models["chatgpt_sol"];
1134        assert_eq!(entry.billing, Some(BillingMode::Subscription));
1135        assert_eq!(entry.catalog_verified, Some(true));
1136        assert!(entry.secret.is_none());
1137        let out = serde_yaml_ng::to_string(&reg).unwrap();
1138        assert!(out.contains("billing: subscription"), "{out}");
1139        assert!(out.contains("catalog_verified: true"), "{out}");
1140    }
1141
1142    /// Entries written before billing metadata existed keep loading and
1143    /// stay unknown — never inheriting a billing mode on reserialize.
1144    #[test]
1145    fn entry_without_billing_metadata_stays_unknown() {
1146        let yaml = r#"schema_version: 1
1147models:
1148  gpt:
1149    provider: openai
1150    model: gpt-4o
1151    secret: env:OPENAI_API_KEY
1152"#;
1153        let reg: ModelRegistry = serde_yaml_ng::from_str(yaml).unwrap();
1154        let entry = &reg.models["gpt"];
1155        assert_eq!(entry.billing, None);
1156        assert_eq!(entry.catalog_verified, None);
1157        let out = serde_yaml_ng::to_string(&reg).unwrap();
1158        assert!(!out.contains("billing"), "{out}");
1159        assert!(!out.contains("catalog_verified"), "{out}");
1160        for (raw, mode) in [
1161            ("subscription", BillingMode::Subscription),
1162            ("usage_billed", BillingMode::UsageBilled),
1163            ("local", BillingMode::Local),
1164        ] {
1165            let m: BillingMode = serde_yaml_ng::from_str(raw).unwrap();
1166            assert_eq!(m, mode);
1167        }
1168    }
1169}