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}
24
25#[derive(Debug, Clone, Default)]
26pub struct AliasEntry {
27 pub model: String,
28}
29
30#[derive(Debug, Clone, Default)]
31pub struct ModelConfig {
32 pub models: HashMap<String, ModelEntry>,
33 pub aliases: HashMap<String, AliasEntry>,
34}
35
36static MODEL_CONFIG: RwLock<Option<ModelConfig>> = RwLock::new(None);
37
38pub fn set_model_config(cfg: ModelConfig) {
39 *MODEL_CONFIG.write().unwrap() = Some(cfg);
40}
41
42pub fn resolve_alias(name: &str) -> String {
43 if let Ok(Some(cfg)) = MODEL_CONFIG.read().as_deref() {
44 if let Some(entry) = cfg.aliases.get(name) {
45 return entry.model.clone();
46 }
47 }
48 name.to_string()
49}
50
51pub fn model_entry(name: &str) -> Option<ModelEntry> {
52 let resolved = resolve_alias(name);
53 if let Ok(Some(cfg)) = MODEL_CONFIG.read().as_deref() {
54 return cfg.models.get(&resolved).cloned();
55 }
56 None
57}
58
59pub fn all_model_entries() -> Vec<(String, ModelEntry)> {
60 if let Ok(Some(cfg)) = MODEL_CONFIG.read().as_deref() {
61 return cfg
62 .models
63 .iter()
64 .map(|(k, v)| (k.clone(), v.clone()))
65 .collect();
66 }
67 Vec::new()
68}
69
70pub fn all_aliases() -> Vec<(String, String)> {
71 if let Ok(Some(cfg)) = MODEL_CONFIG.read().as_deref() {
72 return cfg
73 .aliases
74 .iter()
75 .map(|(k, v)| (k.clone(), v.model.clone()))
76 .collect();
77 }
78 Vec::new()
79}
80
81pub fn model_info(name: &str) -> ModelInfo {
82 let resolved = resolve_alias(name);
83 if let Ok(Some(cfg)) = MODEL_CONFIG.read().as_deref() {
84 if let Some(entry) = cfg.models.get(&resolved) {
85 let (budget, ratio) = builtin_budget(&resolved);
86 return ModelInfo {
87 name: resolved.clone(),
88 context_budget: entry.context_budget.unwrap_or(budget),
89 compact_threshold_ratio: entry.compact_threshold_ratio.unwrap_or(ratio),
90 thinking_enabled: entry.thinking.unwrap_or(false),
91 max_output_tokens: entry.max_tokens,
92 };
93 }
94 }
95 let (budget, ratio) = builtin_budget(&resolved);
96 ModelInfo {
97 name: resolved,
98 context_budget: budget,
99 compact_threshold_ratio: ratio,
100 thinking_enabled: false,
101 max_output_tokens: None,
102 }
103}
104
105fn builtin_budget(name: &str) -> (u64, f64) {
106 let bare = match name.split_once('/') {
107 Some((_, rest)) => rest,
108 None => name,
109 };
110 match bare {
111 n if n.starts_with("claude-opus") => (200_000, 0.8),
112 n if n.starts_with("claude-sonnet") => (200_000, 0.8),
113 n if n.starts_with("claude-haiku") => (200_000, 0.8),
114 n if n.starts_with("claude-") => (200_000, 0.8),
115 n if n.starts_with("gpt-5") => (128_000, 0.8),
116 n if n.starts_with("gpt-4o-mini") => (128_000, 0.8),
117 n if n.starts_with("gpt-4o") => (128_000, 0.8),
118 n if n.starts_with("gpt-4-turbo") => (128_000, 0.8),
119 n if n.starts_with("gpt-4") => (32_000, 0.8),
120 n if n.starts_with("gpt-3.5") => (16_000, 0.8),
121 n if n.starts_with("o1") => (128_000, 0.8),
122 n if n.starts_with("o3") => (128_000, 0.8),
123 n if n.starts_with("glm-5") => (128_000, 0.8),
124 n if n.starts_with("glm-4.5") => (128_000, 0.8),
125 n if n.starts_with("glm-4") => (128_000, 0.8),
126 n if n.starts_with("glm-") => (128_000, 0.8),
127 n if n.starts_with("deepseek-v4") => (1_000_000, 0.8),
128 n if n.starts_with("deepseek-v3") => (128_000, 0.8),
129 n if n.starts_with("deepseek-r1") => (128_000, 0.8),
130 n if n.starts_with("deepseek") => (64_000, 0.8),
131 n if n.starts_with("qwen3") => (128_000, 0.8),
132 n if n.starts_with("qwen-max") => (128_000, 0.8),
133 n if n.starts_with("qwen") => (32_000, 0.8),
134 n if n.starts_with("llama") => (8_000, 0.8),
135 _ => (32_000, 0.8),
136 }
137}
138
139impl ModelInfo {
140 pub fn compact_threshold_tokens(&self) -> u64 {
141 let reserved = self.max_output_tokens.unwrap_or(0) as u64;
142 let available = self.context_budget.saturating_sub(reserved);
143 (available as f64 * self.compact_threshold_ratio) as u64
144 }
145
146 pub fn thinking_enabled(&self) -> bool {
147 self.thinking_enabled
148 }
149}
150
151#[cfg(test)]
152mod tests {
153 use super::*;
154
155 #[test]
156 fn claude_opus_returns_200k() {
157 assert_eq!(model_info("claude-opus-4.7").context_budget, 200_000);
158 }
159
160 #[test]
161 fn gpt_4o_returns_128k() {
162 assert_eq!(model_info("gpt-4o-mini").context_budget, 128_000);
163 assert_eq!(model_info("gpt-4o-2024-08-06").context_budget, 128_000);
164 }
165
166 #[test]
167 fn unknown_model_falls_back_to_32k() {
168 assert_eq!(model_info("mystery-model").context_budget, 32_000);
169 assert_eq!(model_info("").context_budget, 32_000);
170 }
171
172 #[test]
173 fn threshold_is_eighty_percent() {
174 let info = model_info("claude-opus-4.7");
175 assert_eq!(info.compact_threshold_tokens(), 160_000);
176 }
177
178 #[test]
179 fn alias_resolves_to_real_model() {
180 let mut cfg = ModelConfig::default();
181 cfg.aliases.insert(
182 "smart".into(),
183 AliasEntry {
184 model: "claude-opus-4.7".into(),
185 },
186 );
187 set_model_config(cfg);
188 let info = model_info("smart");
189 assert_eq!(info.context_budget, 200_000);
190 assert_eq!(info.name, "claude-opus-4.7");
191 }
192
193 #[test]
194 fn custom_model_overrides_budget() {
195 let mut cfg = ModelConfig::default();
196 cfg.models.insert(
197 "my-local-model".into(),
198 ModelEntry {
199 model: "my-local-model".into(),
200 context_budget: Some(8192),
201 compact_threshold_ratio: Some(0.9),
202 thinking: None,
203 provider: None,
204 api_key: None,
205 base_url: None,
206 max_tokens: None,
207 },
208 );
209 set_model_config(cfg);
210 let info = model_info("my-local-model");
211 assert_eq!(info.context_budget, 8192);
212 assert_eq!(info.compact_threshold_ratio, 0.9);
213 }
214
215 #[test]
216 fn compact_threshold_reserves_configured_output_tokens() {
217 let mut cfg = ModelConfig::default();
218 cfg.models.insert(
219 "large-output".into(),
220 ModelEntry {
221 model: "large-output".into(),
222 context_budget: Some(1_000_000),
223 compact_threshold_ratio: Some(0.8),
224 thinking: None,
225 provider: None,
226 api_key: None,
227 base_url: None,
228 max_tokens: Some(400_000),
229 },
230 );
231 set_model_config(cfg);
232 let info = model_info("large-output");
233 assert_eq!(info.compact_threshold_tokens(), 480_000);
234 }
235
236 #[test]
237 fn alias_chains_through_custom_model() {
238 let mut cfg = ModelConfig::default();
239 cfg.aliases.insert(
240 "default".into(),
241 AliasEntry {
242 model: "my-model".into(),
243 },
244 );
245 cfg.models.insert(
246 "my-model".into(),
247 ModelEntry {
248 model: "my-model".into(),
249 context_budget: Some(65_536),
250 compact_threshold_ratio: None,
251 thinking: None,
252 provider: None,
253 api_key: None,
254 base_url: None,
255 max_tokens: None,
256 },
257 );
258 set_model_config(cfg);
259 let info = model_info("default");
260 assert_eq!(info.name, "my-model");
261 assert_eq!(info.context_budget, 65_536);
262 }
263}