atman_runtime/
model_registry.rs

1use std::collections::HashMap;
2use std::sync::RwLock;
3
4#[derive(Debug, Clone)]
5pub struct ModelInfo {
6    pub name: String,
7    pub context_budget: u64,
8    pub compact_threshold_ratio: f64,
9    pub thinking_enabled: bool,
10    pub max_output_tokens: Option<u32>,
11}
12
13#[derive(Debug, Clone, Default)]
14pub struct ModelEntry {
15    pub model: String,
16    pub provider: Option<String>,
17    pub api_key: Option<String>,
18    pub base_url: Option<String>,
19    pub context_budget: Option<u64>,
20    pub compact_threshold_ratio: Option<f64>,
21    pub thinking: Option<bool>,
22    pub max_tokens: Option<u32>,
23    pub enabled: Option<bool>,
24    #[allow(dead_code)]
25    pub discovered: bool,
26}
27
28#[derive(Debug, Clone, Default)]
29pub struct AliasEntry {
30    pub model: String,
31}
32
33#[derive(Debug, Clone, Default)]
34pub struct ModelConfig {
35    pub models: HashMap<String, ModelEntry>,
36    pub aliases: HashMap<String, AliasEntry>,
37}
38
39static MODEL_CONFIG: RwLock<Option<ModelConfig>> = RwLock::new(None);
40
41/// Model IDs discovered from OAuth providers (e.g. Codex).
42static DISCOVERED_MODELS: RwLock<Vec<String>> = RwLock::new(Vec::new());
43
44pub fn set_discovered_models(models: Vec<String>) {
45    *DISCOVERED_MODELS.write().unwrap() = models;
46}
47
48pub fn discovered_models() -> Vec<String> {
49    DISCOVERED_MODELS.read().unwrap().clone()
50}
51
52/// Set the base model configuration (from config.toml).
53/// Preserves previously registered discovered models.
54pub fn set_model_config(mut cfg: ModelConfig) {
55    let mut guard = MODEL_CONFIG.write().unwrap();
56    if let Some(old) = guard.take() {
57        // Preserve discovered models from old config.
58        for (name, entry) in old.models {
59            if entry.discovered {
60                cfg.models.entry(name).or_insert(entry);
61            }
62        }
63        // Preserve aliases from old config that aren't in the new one.
64        // config.toml aliases are authoritative; only preserve
65        // non-config-sourced aliases (currently none exist).
66    }
67    *guard = Some(cfg);
68}
69
70/// Register additional model entries without clobbering existing ones.
71pub fn register_model_entries(entries: Vec<(String, ModelEntry)>) {
72    let mut guard = MODEL_CONFIG.write().unwrap();
73    let mut cfg = guard.take().unwrap_or_default();
74    for (name, entry) in entries {
75        cfg.models.entry(name).or_insert(entry);
76    }
77    *guard = Some(cfg);
78}
79
80/// Build ModelEntry values from discovered models and register them.
81/// Models are keyed as `<provider_name>:<slug>` (e.g. `Codex:codex/gpt-5.5`).
82pub fn register_discovered(
83    _provider_id: &str,
84    provider_name: &str,
85    models: &[crate::provider::DiscoveredModel],
86) {
87    let entries: Vec<(String, ModelEntry)> = models
88        .iter()
89        .map(|m| {
90            let name = format!("{provider_name}:{}", m.slug);
91            let entry = ModelEntry {
92                model: name.clone(),
93                provider: Some(provider_name.to_string()),
94                context_budget: m.context_budget,
95                thinking: Some(m.thinking),
96                enabled: None,
97                discovered: true,
98                ..Default::default()
99            };
100            (name, entry)
101        })
102        .collect();
103    let slugs: Vec<String> = entries.iter().map(|(name, _)| name.clone()).collect();
104    register_model_entries(entries);
105    set_discovered_models(slugs);
106}
107
108#[derive(Debug, Clone)]
109pub struct ModelRow {
110    pub slug: String,
111    pub provider_name: String,
112    pub context_budget: u64,
113    pub max_output_tokens: Option<u32>,
114    pub thinking: bool,
115}
116
117#[derive(Debug, Clone)]
118pub struct ProviderGroup {
119    pub provider_name: String,
120    pub models: Vec<ModelRow>,
121}
122
123/// Return all models grouped by provider, with complete metadata.
124/// Single canonical source for UI — no manual union of all_model_entries +
125/// discovered_models + aliases.
126pub fn all_provider_groups() -> Vec<ProviderGroup> {
127    let entries = all_model_entries();
128    let mut groups: std::collections::BTreeMap<String, Vec<ModelRow>> =
129        std::collections::BTreeMap::new();
130    for (name, entry) in entries {
131        let info = model_info(&name);
132        let provider = entry.provider.unwrap_or_else(|| "unknown".to_string());
133        let row = ModelRow {
134            slug: name,
135            provider_name: provider.clone(),
136            context_budget: info.context_budget,
137            max_output_tokens: info.max_output_tokens,
138            thinking: info.thinking_enabled(),
139        };
140        groups.entry(provider).or_default().push(row);
141    }
142    groups
143        .into_iter()
144        .map(|(provider_name, models)| ProviderGroup {
145            provider_name,
146            models,
147        })
148        .collect()
149}
150
151pub fn resolve_alias(name: &str) -> String {
152    if let Ok(Some(cfg)) = MODEL_CONFIG.read().as_deref() {
153        if let Some(entry) = cfg.aliases.get(name) {
154            return entry.model.clone();
155        }
156    }
157    name.to_string()
158}
159
160pub fn model_entry(name: &str) -> Option<ModelEntry> {
161    let resolved = resolve_alias(name);
162    if let Ok(Some(cfg)) = MODEL_CONFIG.read().as_deref() {
163        return cfg.models.get(&resolved).cloned();
164    }
165    None
166}
167
168pub fn all_model_entries() -> Vec<(String, ModelEntry)> {
169    if let Ok(Some(cfg)) = MODEL_CONFIG.read().as_deref() {
170        return cfg
171            .models
172            .iter()
173            .map(|(k, v)| (k.clone(), v.clone()))
174            .collect();
175    }
176    Vec::new()
177}
178
179pub fn all_aliases() -> Vec<(String, String)> {
180    if let Ok(Some(cfg)) = MODEL_CONFIG.read().as_deref() {
181        return cfg
182            .aliases
183            .iter()
184            .map(|(k, v)| (k.clone(), v.model.clone()))
185            .collect();
186    }
187    Vec::new()
188}
189
190pub fn model_info(name: &str) -> ModelInfo {
191    let resolved = resolve_alias(name);
192    if let Ok(Some(cfg)) = MODEL_CONFIG.read().as_deref() {
193        if let Some(entry) = cfg.models.get(&resolved) {
194            return ModelInfo {
195                name: resolved.clone(),
196                context_budget: entry.context_budget.unwrap_or(0),
197                compact_threshold_ratio: entry.compact_threshold_ratio.unwrap_or(0.8),
198                thinking_enabled: entry.thinking.unwrap_or(false),
199                max_output_tokens: entry.max_tokens,
200            };
201        }
202    }
203    ModelInfo {
204        name: resolved,
205        context_budget: 0,
206        compact_threshold_ratio: 0.8,
207        thinking_enabled: false,
208        max_output_tokens: None,
209    }
210}
211
212impl ModelInfo {
213    pub fn compact_threshold_tokens(&self) -> u64 {
214        let reserved = self.max_output_tokens.unwrap_or(0) as u64;
215        let available = self.context_budget.saturating_sub(reserved);
216        (available as f64 * self.compact_threshold_ratio) as u64
217    }
218
219    pub fn compaction_trigger_threshold(&self) -> u64 {
220        let budget = self.context_budget;
221
222        let configured_output = self.max_output_tokens.unwrap_or(32_000) as u64;
223        let output_cap = (budget as f64 * 0.20) as u64;
224        let output_reserve = configured_output.min(output_cap).max(8_000);
225
226        let safety = (budget as f64 * 0.05) as u64;
227        let safety_margin = safety.max(4_000);
228
229        let trigger = budget
230            .saturating_sub(output_reserve)
231            .saturating_sub(safety_margin);
232        let floor = (budget as f64 * 0.50) as u64;
233        let ceiling = (budget as f64 * 0.95) as u64;
234        trigger.clamp(floor, ceiling)
235    }
236
237    pub fn compaction_target_after(&self) -> u64 {
238        let trigger = self.compaction_trigger_threshold();
239        let budget_cap = (self.context_budget as f64 * 0.60) as u64;
240        let trigger_cap = (trigger as f64 * 0.75) as u64;
241        budget_cap.min(trigger_cap)
242    }
243
244    pub fn thinking_enabled(&self) -> bool {
245        self.thinking_enabled
246    }
247}
248
249// ── Alias CRUD (writes config.toml) ──
250
251fn read_config_toml() -> Option<String> {
252    let path = crate::storage::config_dir().ok()?.join("config.toml");
253    std::fs::read_to_string(&path).ok()
254}
255
256fn write_config_toml(text: &str) -> anyhow::Result<()> {
257    let dir = crate::storage::config_dir().map_err(|e| anyhow::anyhow!("config dir: {e}"))?;
258    std::fs::create_dir_all(&dir)?;
259    let path = dir.join("config.toml");
260    std::fs::write(&path, text)?;
261    Ok(())
262}
263
264fn reload_from_text(text: &str) {
265    let Ok(raw) = toml::from_str::<toml::Value>(text) else {
266        return;
267    };
268    let mut guard = MODEL_CONFIG.write().unwrap();
269    let mut cfg = guard.take().unwrap_or_default();
270
271    // Update aliases from config.toml — preserve all other state
272    // (discovered models, config-defined models).
273    cfg.aliases.clear();
274    if let Some(aliases) = raw.get("alias").and_then(|a| a.as_table()) {
275        for (name, entry) in aliases {
276            if let Some(model) = entry.get("model").and_then(|m| m.as_str()) {
277                cfg.aliases.insert(
278                    name.clone(),
279                    AliasEntry {
280                        model: model.to_string(),
281                    },
282                );
283            }
284        }
285    }
286
287    // Update config-defined models — only update existing keys or add
288    // new ones; never remove entries that aren't in config.toml.
289    if let Some(models) = raw.get("models").and_then(|m| m.as_table()) {
290        for (name, entry) in models {
291            let provider = entry
292                .get("provider")
293                .and_then(|v| v.as_str())
294                .map(String::from);
295            let api_key = entry
296                .get("api_key")
297                .and_then(|v| v.as_str())
298                .map(String::from);
299            let base_url = entry
300                .get("base_url")
301                .and_then(|v| v.as_str())
302                .map(String::from);
303            let context_budget = entry
304                .get("context_budget")
305                .and_then(|v| v.as_integer())
306                .map(|n| n as u64);
307            let thinking = entry.get("thinking").and_then(|v| v.as_bool());
308            let max_tokens = entry
309                .get("max_tokens")
310                .and_then(|v| v.as_integer())
311                .map(|n| n as u32);
312            let model = entry
313                .get("model")
314                .and_then(|v| v.as_str())
315                .map(String::from);
316            cfg.models.insert(
317                name.clone(),
318                ModelEntry {
319                    model: model.unwrap_or_default(),
320                    provider,
321                    api_key,
322                    base_url,
323                    context_budget,
324                    compact_threshold_ratio: None,
325                    thinking,
326                    max_tokens,
327                    enabled: None,
328                    discovered: false,
329                },
330            );
331        }
332    }
333
334    *guard = Some(cfg);
335}
336
337pub fn add_alias_to_config(alias: &str, model: &str) -> anyhow::Result<()> {
338    let text = read_config_toml().unwrap_or_default();
339    let mut raw: toml::Value = if text.trim().is_empty() {
340        toml::Value::Table(toml::value::Table::new())
341    } else {
342        toml::from_str(&text).map_err(|e| anyhow::anyhow!("parse config.toml: {e}"))?
343    };
344    let aliases = raw
345        .as_table_mut()
346        .ok_or_else(|| anyhow::anyhow!("config.toml is not a table"))?
347        .entry("alias")
348        .or_insert_with(|| toml::Value::Table(toml::value::Table::new()));
349    if let Some(table) = aliases.as_table_mut() {
350        let mut entry = toml::value::Table::new();
351        entry.insert("model".to_string(), toml::Value::String(model.to_string()));
352        table.insert(alias.to_string(), toml::Value::Table(entry));
353    }
354    let new_text = toml::to_string_pretty(&raw).map_err(|e| anyhow::anyhow!("serialize: {e}"))?;
355    write_config_toml(&new_text)?;
356    reload_from_text(&new_text);
357    Ok(())
358}
359
360pub fn remove_alias_from_config(alias: &str) -> anyhow::Result<()> {
361    let text = read_config_toml().unwrap_or_default();
362    let mut raw: toml::Value = toml::from_str(&text).map_err(|e| anyhow::anyhow!("parse: {e}"))?;
363    if let Some(table) = raw.get_mut("alias").and_then(|a| a.as_table_mut()) {
364        table.remove(alias);
365    }
366    let new_text = toml::to_string_pretty(&raw).map_err(|e| anyhow::anyhow!("serialize: {e}"))?;
367    write_config_toml(&new_text)?;
368    reload_from_text(&new_text);
369    Ok(())
370}
371
372pub fn update_alias_in_config(
373    old_alias: &str,
374    new_alias: &str,
375    new_model: &str,
376) -> anyhow::Result<()> {
377    let text = read_config_toml().unwrap_or_default();
378    let mut raw: toml::Value = toml::from_str(&text).map_err(|e| anyhow::anyhow!("parse: {e}"))?;
379    if let Some(table) = raw.get_mut("alias").and_then(|a| a.as_table_mut()) {
380        table.remove(old_alias);
381        let mut entry = toml::value::Table::new();
382        entry.insert(
383            "model".to_string(),
384            toml::Value::String(new_model.to_string()),
385        );
386        table.insert(new_alias.to_string(), toml::Value::Table(entry));
387    }
388    let new_text = toml::to_string_pretty(&raw).map_err(|e| anyhow::anyhow!("serialize: {e}"))?;
389    write_config_toml(&new_text)?;
390    reload_from_text(&new_text);
391    Ok(())
392}
393
394#[cfg(test)]
395mod tests {
396    use super::*;
397    use std::sync::Mutex as StdMutex;
398
399    /// Tests that mutate MODEL_CONFIG must hold this lock to avoid races
400    /// when cargo test runs them in parallel.
401    static TEST_CFG_LOCK: StdMutex<()> = StdMutex::new(());
402
403    #[test]
404    fn unregistered_model_returns_zero_budget() {
405        let _lock = TEST_CFG_LOCK.lock().unwrap();
406        *MODEL_CONFIG.write().unwrap() = None;
407        assert_eq!(model_info("mystery-model").context_budget, 0);
408        assert_eq!(model_info("").context_budget, 0);
409    }
410
411    #[test]
412    fn threshold_is_eighty_percent() {
413        let _lock = TEST_CFG_LOCK.lock().unwrap();
414        let mut cfg = ModelConfig::default();
415        cfg.models.insert(
416            "claude-opus-4.7".into(),
417            ModelEntry {
418                model: "claude-opus-4.7".into(),
419                context_budget: Some(200_000),
420                compact_threshold_ratio: Some(0.8),
421                thinking: None,
422                ..Default::default()
423            },
424        );
425        set_model_config(cfg);
426        let info = model_info("claude-opus-4.7");
427        assert_eq!(info.compact_threshold_tokens(), 160_000);
428    }
429
430    #[test]
431    fn compaction_trigger_is_near_budget_top() {
432        let _lock = TEST_CFG_LOCK.lock().unwrap();
433        let mut cfg = ModelConfig::default();
434        cfg.models.insert(
435            "claude-opus-4.7".into(),
436            ModelEntry {
437                model: "claude-opus-4.7".into(),
438                context_budget: Some(200_000),
439                compact_threshold_ratio: Some(0.8),
440                thinking: None,
441                ..Default::default()
442            },
443        );
444        set_model_config(cfg);
445        let info = model_info("claude-opus-4.7");
446        let trigger = info.compaction_trigger_threshold();
447        assert!(
448            trigger > 150_000 && trigger <= 190_000,
449            "trigger should be near the top of the budget, got {trigger}"
450        );
451    }
452
453    #[test]
454    fn compaction_target_is_lower_than_trigger() {
455        let _lock = TEST_CFG_LOCK.lock().unwrap();
456        let mut cfg = ModelConfig::default();
457        cfg.models.insert(
458            "claude-opus-4.7".into(),
459            ModelEntry {
460                model: "claude-opus-4.7".into(),
461                context_budget: Some(200_000),
462                compact_threshold_ratio: Some(0.8),
463                thinking: None,
464                ..Default::default()
465            },
466        );
467        set_model_config(cfg);
468        let info = model_info("claude-opus-4.7");
469        let trigger = info.compaction_trigger_threshold();
470        let target = info.compaction_target_after();
471        assert!(
472            target < trigger,
473            "target {target} should be less than trigger {trigger}"
474        );
475    }
476
477    #[test]
478    fn alias_resolves_to_real_model() {
479        let _lock = TEST_CFG_LOCK.lock().unwrap();
480        let mut cfg = ModelConfig::default();
481        cfg.models.insert(
482            "claude-opus-4.7".into(),
483            ModelEntry {
484                model: "claude-opus-4.7".into(),
485                context_budget: Some(200_000),
486                ..Default::default()
487            },
488        );
489        cfg.aliases.insert(
490            "smart".into(),
491            AliasEntry {
492                model: "claude-opus-4.7".into(),
493            },
494        );
495        set_model_config(cfg);
496        let info = model_info("smart");
497        assert_eq!(info.context_budget, 200_000);
498        assert_eq!(info.name, "claude-opus-4.7");
499    }
500
501    #[test]
502    fn custom_model_overrides_budget() {
503        let _lock = TEST_CFG_LOCK.lock().unwrap();
504        let mut cfg = ModelConfig::default();
505        cfg.models.insert(
506            "my-local-model".into(),
507            ModelEntry {
508                model: "my-local-model".into(),
509                context_budget: Some(8192),
510                compact_threshold_ratio: Some(0.9),
511                thinking: None,
512                ..Default::default()
513            },
514        );
515        set_model_config(cfg);
516        let info = model_info("my-local-model");
517        assert_eq!(info.context_budget, 8192);
518        assert_eq!(info.compact_threshold_ratio, 0.9);
519    }
520
521    #[test]
522    fn compact_threshold_reserves_configured_output_tokens() {
523        let _lock = TEST_CFG_LOCK.lock().unwrap();
524        let mut cfg = ModelConfig::default();
525        cfg.models.insert(
526            "large-output".into(),
527            ModelEntry {
528                model: "large-output".into(),
529                context_budget: Some(1_000_000),
530                compact_threshold_ratio: Some(0.8),
531                thinking: None,
532                max_tokens: Some(400_000),
533                ..Default::default()
534            },
535        );
536        set_model_config(cfg);
537        let info = model_info("large-output");
538        assert_eq!(info.compact_threshold_tokens(), 480_000);
539        let trigger = info.compaction_trigger_threshold();
540        assert!(
541            trigger > 700_000,
542            "trigger with capped output reserve should be > 700K, got {trigger}"
543        );
544    }
545
546    #[test]
547    fn alias_chains_through_custom_model() {
548        let _lock = TEST_CFG_LOCK.lock().unwrap();
549        let mut cfg = ModelConfig::default();
550        cfg.aliases.insert(
551            "default".into(),
552            AliasEntry {
553                model: "my-model".into(),
554            },
555        );
556        cfg.models.insert(
557            "my-model".into(),
558            ModelEntry {
559                model: "my-model".into(),
560                context_budget: Some(65_536),
561                compact_threshold_ratio: None,
562                thinking: None,
563                ..Default::default()
564            },
565        );
566        set_model_config(cfg);
567        let info = model_info("default");
568        assert_eq!(info.name, "my-model");
569        assert_eq!(info.context_budget, 65_536);
570    }
571
572    #[test]
573    fn discovered_models_survive_set_model_config() {
574        let _lock = TEST_CFG_LOCK.lock().unwrap();
575        // Register discovered models first.
576        register_discovered(
577            "pid-abc",
578            "Codex",
579            &[crate::provider::DiscoveredModel {
580                slug: "codex/gpt-5".to_string(),
581                context_budget: Some(128_000),
582                thinking: true,
583            }],
584        );
585        assert!(model_entry("Codex:codex/gpt-5").is_some());
586
587        // Simulate config reload from config.toml.
588        let mut cfg = ModelConfig::default();
589        cfg.aliases.insert(
590            "cheap".into(),
591            AliasEntry {
592                model: "claude-opus-4.7".into(),
593            },
594        );
595        set_model_config(cfg);
596
597        // Discovered models should still be there.
598        assert!(
599            model_entry("Codex:codex/gpt-5").is_some(),
600            "discovered models should survive set_model_config"
601        );
602        // Alias from config should work.
603        assert_eq!(resolve_alias("cheap"), "claude-opus-4.7");
604    }
605
606    #[test]
607    fn discovered_models_survive_reload_from_text_alias_crud() {
608        let _lock = TEST_CFG_LOCK.lock().unwrap();
609        register_discovered(
610            "pid-abc",
611            "Codex",
612            &[crate::provider::DiscoveredModel {
613                slug: "codex/gpt-5".to_string(),
614                context_budget: Some(128_000),
615                thinking: true,
616            }],
617        );
618
619        // Simulate alias add → reload_from_text
620        let toml = r#"
621[alias]
622smart = { model = "Codex:codex/gpt-5" }
623"#;
624        reload_from_text(toml);
625
626        assert!(
627            model_entry("Codex:codex/gpt-5").is_some(),
628            "discovered models should survive alias CRUD"
629        );
630        assert_eq!(resolve_alias("smart"), "Codex:codex/gpt-5");
631    }
632}
atman_runtime/model_registry.rs

atman_runtime/
model_registry.rs
