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    /// True for models registered via discover (OAuth providers),
24    /// false for config.toml-defined models.
25    #[allow(dead_code)]
26    pub discovered: bool,
27}
28
29#[derive(Debug, Clone, Default)]
30pub struct AliasEntry {
31    pub model: String,
32}
33
34#[derive(Debug, Clone, Default)]
35pub struct ModelConfig {
36    pub models: HashMap<String, ModelEntry>,
37    pub aliases: HashMap<String, AliasEntry>,
38}
39
40static MODEL_CONFIG: RwLock<Option<ModelConfig>> = RwLock::new(None);
41
42/// Model IDs discovered from OAuth providers (e.g. Codex).
43static DISCOVERED_MODELS: RwLock<Vec<String>> = RwLock::new(Vec::new());
44
45pub fn set_discovered_models(models: Vec<String>) {
46    *DISCOVERED_MODELS.write().unwrap() = models;
47}
48
49pub fn discovered_models() -> Vec<String> {
50    DISCOVERED_MODELS.read().unwrap().clone()
51}
52
53/// Set the base model configuration (from config.toml).
54/// Preserves previously registered discovered models.
55pub fn set_model_config(mut cfg: ModelConfig) {
56    let mut guard = MODEL_CONFIG.write().unwrap();
57    if let Some(old) = guard.take() {
58        // Preserve discovered models from old config.
59        for (name, entry) in old.models {
60            if entry.discovered {
61                cfg.models.entry(name).or_insert(entry);
62            }
63        }
64        // Preserve aliases from old config that aren't in the new one.
65        // config.toml aliases are authoritative; only preserve
66        // non-config-sourced aliases (currently none exist).
67    }
68    *guard = Some(cfg);
69}
70
71/// Register additional model entries without clobbering existing ones.
72pub fn register_model_entries(entries: Vec<(String, ModelEntry)>) {
73    let mut guard = MODEL_CONFIG.write().unwrap();
74    let mut cfg = guard.take().unwrap_or_default();
75    for (name, entry) in entries {
76        cfg.models.entry(name).or_insert(entry);
77    }
78    *guard = Some(cfg);
79}
80
81/// Build ModelEntry values from discovered models and register them.
82/// Models are keyed as `<provider_name>:<slug>` (e.g. `Codex:codex/gpt-5.5`).
83pub fn register_discovered(
84    _provider_id: &str,
85    provider_name: &str,
86    models: &[crate::provider::DiscoveredModel],
87) {
88    let entries: Vec<(String, ModelEntry)> = models
89        .iter()
90        .map(|m| {
91            let name = format!("{provider_name}:{}", m.slug);
92            let entry = ModelEntry {
93                model: name.clone(),
94                provider: Some(provider_name.to_string()),
95                context_budget: m.context_budget,
96                thinking: Some(m.thinking),
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                    discovered: false,
328                },
329            );
330        }
331    }
332
333    *guard = Some(cfg);
334}
335
336pub fn add_alias_to_config(alias: &str, model: &str) -> anyhow::Result<()> {
337    let text = read_config_toml().unwrap_or_default();
338    let mut raw: toml::Value = if text.trim().is_empty() {
339        toml::Value::Table(toml::value::Table::new())
340    } else {
341        toml::from_str(&text).map_err(|e| anyhow::anyhow!("parse config.toml: {e}"))?
342    };
343    let aliases = raw
344        .as_table_mut()
345        .ok_or_else(|| anyhow::anyhow!("config.toml is not a table"))?
346        .entry("alias")
347        .or_insert_with(|| toml::Value::Table(toml::value::Table::new()));
348    if let Some(table) = aliases.as_table_mut() {
349        let mut entry = toml::value::Table::new();
350        entry.insert("model".to_string(), toml::Value::String(model.to_string()));
351        table.insert(alias.to_string(), toml::Value::Table(entry));
352    }
353    let new_text = toml::to_string_pretty(&raw).map_err(|e| anyhow::anyhow!("serialize: {e}"))?;
354    write_config_toml(&new_text)?;
355    reload_from_text(&new_text);
356    Ok(())
357}
358
359pub fn remove_alias_from_config(alias: &str) -> anyhow::Result<()> {
360    let text = read_config_toml().unwrap_or_default();
361    let mut raw: toml::Value = toml::from_str(&text).map_err(|e| anyhow::anyhow!("parse: {e}"))?;
362    if let Some(table) = raw.get_mut("alias").and_then(|a| a.as_table_mut()) {
363        table.remove(alias);
364    }
365    let new_text = toml::to_string_pretty(&raw).map_err(|e| anyhow::anyhow!("serialize: {e}"))?;
366    write_config_toml(&new_text)?;
367    reload_from_text(&new_text);
368    Ok(())
369}
370
371pub fn update_alias_in_config(
372    old_alias: &str,
373    new_alias: &str,
374    new_model: &str,
375) -> anyhow::Result<()> {
376    let text = read_config_toml().unwrap_or_default();
377    let mut raw: toml::Value = toml::from_str(&text).map_err(|e| anyhow::anyhow!("parse: {e}"))?;
378    if let Some(table) = raw.get_mut("alias").and_then(|a| a.as_table_mut()) {
379        table.remove(old_alias);
380        let mut entry = toml::value::Table::new();
381        entry.insert(
382            "model".to_string(),
383            toml::Value::String(new_model.to_string()),
384        );
385        table.insert(new_alias.to_string(), toml::Value::Table(entry));
386    }
387    let new_text = toml::to_string_pretty(&raw).map_err(|e| anyhow::anyhow!("serialize: {e}"))?;
388    write_config_toml(&new_text)?;
389    reload_from_text(&new_text);
390    Ok(())
391}
392
393#[cfg(test)]
394mod tests {
395    use super::*;
396    use std::sync::Mutex as StdMutex;
397
398    /// Tests that mutate MODEL_CONFIG must hold this lock to avoid races
399    /// when cargo test runs them in parallel.
400    static TEST_CFG_LOCK: StdMutex<()> = StdMutex::new(());
401
402    #[test]
403    fn unregistered_model_returns_zero_budget() {
404        let _lock = TEST_CFG_LOCK.lock().unwrap();
405        *MODEL_CONFIG.write().unwrap() = None;
406        assert_eq!(model_info("mystery-model").context_budget, 0);
407        assert_eq!(model_info("").context_budget, 0);
408    }
409
410    #[test]
411    fn threshold_is_eighty_percent() {
412        let _lock = TEST_CFG_LOCK.lock().unwrap();
413        let mut cfg = ModelConfig::default();
414        cfg.models.insert(
415            "claude-opus-4.7".into(),
416            ModelEntry {
417                model: "claude-opus-4.7".into(),
418                context_budget: Some(200_000),
419                compact_threshold_ratio: Some(0.8),
420                thinking: None,
421                ..Default::default()
422            },
423        );
424        set_model_config(cfg);
425        let info = model_info("claude-opus-4.7");
426        assert_eq!(info.compact_threshold_tokens(), 160_000);
427    }
428
429    #[test]
430    fn compaction_trigger_is_near_budget_top() {
431        let _lock = TEST_CFG_LOCK.lock().unwrap();
432        let mut cfg = ModelConfig::default();
433        cfg.models.insert(
434            "claude-opus-4.7".into(),
435            ModelEntry {
436                model: "claude-opus-4.7".into(),
437                context_budget: Some(200_000),
438                compact_threshold_ratio: Some(0.8),
439                thinking: None,
440                ..Default::default()
441            },
442        );
443        set_model_config(cfg);
444        let info = model_info("claude-opus-4.7");
445        let trigger = info.compaction_trigger_threshold();
446        assert!(
447            trigger > 150_000 && trigger <= 190_000,
448            "trigger should be near the top of the budget, got {trigger}"
449        );
450    }
451
452    #[test]
453    fn compaction_target_is_lower_than_trigger() {
454        let _lock = TEST_CFG_LOCK.lock().unwrap();
455        let mut cfg = ModelConfig::default();
456        cfg.models.insert(
457            "claude-opus-4.7".into(),
458            ModelEntry {
459                model: "claude-opus-4.7".into(),
460                context_budget: Some(200_000),
461                compact_threshold_ratio: Some(0.8),
462                thinking: None,
463                ..Default::default()
464            },
465        );
466        set_model_config(cfg);
467        let info = model_info("claude-opus-4.7");
468        let trigger = info.compaction_trigger_threshold();
469        let target = info.compaction_target_after();
470        assert!(
471            target < trigger,
472            "target {target} should be less than trigger {trigger}"
473        );
474    }
475
476    #[test]
477    fn alias_resolves_to_real_model() {
478        let _lock = TEST_CFG_LOCK.lock().unwrap();
479        let mut cfg = ModelConfig::default();
480        cfg.models.insert(
481            "claude-opus-4.7".into(),
482            ModelEntry {
483                model: "claude-opus-4.7".into(),
484                context_budget: Some(200_000),
485                ..Default::default()
486            },
487        );
488        cfg.aliases.insert(
489            "smart".into(),
490            AliasEntry {
491                model: "claude-opus-4.7".into(),
492            },
493        );
494        set_model_config(cfg);
495        let info = model_info("smart");
496        assert_eq!(info.context_budget, 200_000);
497        assert_eq!(info.name, "claude-opus-4.7");
498    }
499
500    #[test]
501    fn custom_model_overrides_budget() {
502        let _lock = TEST_CFG_LOCK.lock().unwrap();
503        let mut cfg = ModelConfig::default();
504        cfg.models.insert(
505            "my-local-model".into(),
506            ModelEntry {
507                model: "my-local-model".into(),
508                context_budget: Some(8192),
509                compact_threshold_ratio: Some(0.9),
510                thinking: None,
511                ..Default::default()
512            },
513        );
514        set_model_config(cfg);
515        let info = model_info("my-local-model");
516        assert_eq!(info.context_budget, 8192);
517        assert_eq!(info.compact_threshold_ratio, 0.9);
518    }
519
520    #[test]
521    fn compact_threshold_reserves_configured_output_tokens() {
522        let _lock = TEST_CFG_LOCK.lock().unwrap();
523        let mut cfg = ModelConfig::default();
524        cfg.models.insert(
525            "large-output".into(),
526            ModelEntry {
527                model: "large-output".into(),
528                context_budget: Some(1_000_000),
529                compact_threshold_ratio: Some(0.8),
530                thinking: None,
531                max_tokens: Some(400_000),
532                ..Default::default()
533            },
534        );
535        set_model_config(cfg);
536        let info = model_info("large-output");
537        assert_eq!(info.compact_threshold_tokens(), 480_000);
538        let trigger = info.compaction_trigger_threshold();
539        assert!(
540            trigger > 700_000,
541            "trigger with capped output reserve should be > 700K, got {trigger}"
542        );
543    }
544
545    #[test]
546    fn alias_chains_through_custom_model() {
547        let _lock = TEST_CFG_LOCK.lock().unwrap();
548        let mut cfg = ModelConfig::default();
549        cfg.aliases.insert(
550            "default".into(),
551            AliasEntry {
552                model: "my-model".into(),
553            },
554        );
555        cfg.models.insert(
556            "my-model".into(),
557            ModelEntry {
558                model: "my-model".into(),
559                context_budget: Some(65_536),
560                compact_threshold_ratio: None,
561                thinking: None,
562                ..Default::default()
563            },
564        );
565        set_model_config(cfg);
566        let info = model_info("default");
567        assert_eq!(info.name, "my-model");
568        assert_eq!(info.context_budget, 65_536);
569    }
570
571    #[test]
572    fn discovered_models_survive_set_model_config() {
573        let _lock = TEST_CFG_LOCK.lock().unwrap();
574        // Register discovered models first.
575        register_discovered(
576            "pid-abc",
577            "Codex",
578            &[crate::provider::DiscoveredModel {
579                slug: "codex/gpt-5".to_string(),
580                context_budget: Some(128_000),
581                thinking: true,
582            }],
583        );
584        assert!(model_entry("Codex:codex/gpt-5").is_some());
585
586        // Simulate config reload from config.toml.
587        let mut cfg = ModelConfig::default();
588        cfg.aliases.insert(
589            "cheap".into(),
590            AliasEntry {
591                model: "claude-opus-4.7".into(),
592            },
593        );
594        set_model_config(cfg);
595
596        // Discovered models should still be there.
597        assert!(
598            model_entry("Codex:codex/gpt-5").is_some(),
599            "discovered models should survive set_model_config"
600        );
601        // Alias from config should work.
602        assert_eq!(resolve_alias("cheap"), "claude-opus-4.7");
603    }
604
605    #[test]
606    fn discovered_models_survive_reload_from_text_alias_crud() {
607        let _lock = TEST_CFG_LOCK.lock().unwrap();
608        register_discovered(
609            "pid-abc",
610            "Codex",
611            &[crate::provider::DiscoveredModel {
612                slug: "codex/gpt-5".to_string(),
613                context_budget: Some(128_000),
614                thinking: true,
615            }],
616        );
617
618        // Simulate alias add → reload_from_text
619        let toml = r#"
620[alias]
621smart = { model = "Codex:codex/gpt-5" }
622"#;
623        reload_from_text(toml);
624
625        assert!(
626            model_entry("Codex:codex/gpt-5").is_some(),
627            "discovered models should survive alias CRUD"
628        );
629        assert_eq!(resolve_alias("smart"), "Codex:codex/gpt-5");
630    }
631}
atman_runtime/model_registry.rs

atman_runtime/
model_registry.rs
