1use std::path::{Path, PathBuf};
2
3use anyhow::{bail, Context, Result};
4use toml_edit::{value, Item, Table};
5
6use super::{
7 config_path, ensure_config_defaults, host_runtime_config_from_doc, normalize_host,
8 profile_from_doc, read_config_doc_or_default, write_config_doc, MemoryAiExecutor, CLAUDE_HOST,
9 CODEX_HOST, DEFAULT_CODEX_MODEL,
10};
11
12#[derive(Clone, Copy, Debug, PartialEq, Eq)]
13pub struct ModelPreset {
14 pub name: &'static str,
15 pub model: &'static str,
16 pub reasoning_effort: Option<&'static str>,
17 pub description: &'static str,
18}
19
20#[derive(Clone, Debug, PartialEq, Eq)]
21pub struct ModelStatus {
22 pub host: Option<String>,
23 pub profile_name: String,
24 pub executor: MemoryAiExecutor,
25 pub model: String,
26 pub reasoning_effort: Option<String>,
27 pub config_path: PathBuf,
28}
29
30#[derive(Clone, Debug, PartialEq, Eq)]
31pub struct ModelChange {
32 pub host: Option<String>,
33 pub profile_name: String,
34 pub executor: MemoryAiExecutor,
35 pub old_model: String,
36 pub new_model: String,
37 pub old_reasoning_effort: Option<String>,
38 pub new_reasoning_effort: Option<String>,
39 pub config_path: PathBuf,
40 pub backup_path: Option<PathBuf>,
41 pub dry_run: bool,
42}
43
44pub const MODEL_PRESETS: &[ModelPreset] = &[
45 ModelPreset {
46 name: "cheap",
47 model: "gpt-5.4-mini",
48 reasoning_effort: Some("low"),
49 description: "lower-cost Codex mini profile with low reasoning",
50 },
51 ModelPreset {
52 name: "balanced",
53 model: "gpt-5.4-mini",
54 reasoning_effort: Some("medium"),
55 description: "Codex mini profile with more reasoning for extraction quality",
56 },
57 ModelPreset {
58 name: "quality",
59 model: DEFAULT_CODEX_MODEL,
60 reasoning_effort: Some("medium"),
61 description: "higher-quality Codex profile; higher cost",
62 },
63 ModelPreset {
64 name: "auto",
65 model: "auto",
66 reasoning_effort: None,
67 description: "omit --model and let Codex choose its default",
68 },
69];
70
71impl MemoryAiExecutor {
72 pub fn as_str(self) -> &'static str {
73 match self {
74 Self::Http => "http",
75 Self::ClaudeCli => "claude-cli",
76 Self::CodexCli => "codex-cli",
77 }
78 }
79}
80
81pub fn model_statuses() -> Result<Vec<ModelStatus>> {
82 [CODEX_HOST, CLAUDE_HOST]
83 .iter()
84 .map(|host| model_status(Some(host), None))
85 .collect()
86}
87
88pub fn model_status(host: Option<&str>, profile: Option<&str>) -> Result<ModelStatus> {
89 let mut doc = read_config_doc_or_default()?;
90 ensure_config_defaults(&mut doc, &[CLAUDE_HOST, CODEX_HOST])?;
91 let selection = select_profile_from_doc(&doc, host, profile)?;
92 let resolved = profile_from_doc(&doc, &selection.profile_name)?;
93 Ok(ModelStatus {
94 host: selection.host,
95 profile_name: resolved.profile_name,
96 executor: resolved.executor,
97 model: resolved.model.unwrap_or_else(|| "auto".to_string()),
98 reasoning_effort: resolved.reasoning_effort,
99 config_path: config_path()?,
100 })
101}
102
103pub fn set_model(
104 host: Option<&str>,
105 profile: Option<&str>,
106 target: &str,
107 reasoning_effort: Option<&str>,
108 dry_run: bool,
109) -> Result<ModelChange> {
110 let path = config_path()?;
111 let mut doc = read_config_doc_or_default()?;
112 ensure_config_defaults(&mut doc, &[CLAUDE_HOST, CODEX_HOST])?;
113 let selection = select_profile_from_doc(&doc, host, profile)?;
114 let before = profile_from_doc(&doc, &selection.profile_name)?;
115 let target = resolve_model_target(target, reasoning_effort, before.executor)?;
116 let old_model = before.model.clone().unwrap_or_else(|| "auto".to_string());
117 let old_reasoning_effort = before.reasoning_effort.clone();
118 let new_reasoning_effort = if target.update_reasoning {
119 target.reasoning_effort.clone()
120 } else {
121 old_reasoning_effort.clone()
122 };
123 let change = ModelChange {
124 host: selection.host.clone(),
125 profile_name: selection.profile_name.clone(),
126 executor: before.executor,
127 old_model,
128 new_model: target.model.clone(),
129 old_reasoning_effort,
130 new_reasoning_effort: new_reasoning_effort.clone(),
131 config_path: path.clone(),
132 backup_path: (!dry_run).then(|| backup_path_for_config(&path)),
133 dry_run,
134 };
135 if dry_run {
136 return Ok(change);
137 }
138
139 let backup_path = backup_path_for_config(&path);
140 write_config_doc(&backup_path, &doc)?;
141 let profile_table = profile_table_mut(&mut doc, &selection.profile_name)?;
142 profile_table["model"] = value(target.model);
143 match new_reasoning_effort {
144 Some(reasoning) => profile_table["reasoning_effort"] = value(reasoning),
145 None => {
146 profile_table.remove("reasoning_effort");
147 }
148 }
149 write_config_doc(&path, &doc)?;
150 Ok(change)
151}
152
153pub fn rollback_model_config() -> Result<(PathBuf, PathBuf)> {
154 let path = config_path()?;
155 let backup_path = backup_path_for_config(&path);
156 if !backup_path.exists() {
157 bail!(
158 "no model config backup found at {}; run `remem model use ...` first",
159 backup_path.display()
160 );
161 }
162 let backup = std::fs::read_to_string(&backup_path)
163 .with_context(|| format!("read model config backup {}", backup_path.display()))?;
164 crate::atomic_file::write_atomic(&path, backup)
165 .with_context(|| format!("restore {} from {}", path.display(), backup_path.display()))?;
166 Ok((path, backup_path))
167}
168
169struct ProfileSelection {
170 host: Option<String>,
171 profile_name: String,
172}
173
174struct ModelTarget {
175 model: String,
176 reasoning_effort: Option<String>,
177 update_reasoning: bool,
178}
179
180fn select_profile_from_doc(
181 doc: &toml_edit::DocumentMut,
182 host: Option<&str>,
183 profile: Option<&str>,
184) -> Result<ProfileSelection> {
185 if host.is_some() && profile.is_some() {
186 bail!("--host and --profile are mutually exclusive");
187 }
188 if let Some(profile) = profile.map(str::trim).filter(|profile| !profile.is_empty()) {
189 return Ok(ProfileSelection {
190 host: None,
191 profile_name: profile.to_string(),
192 });
193 }
194 let host = host
195 .map(normalize_host)
196 .filter(|host| !host.trim().is_empty())
197 .unwrap_or_else(|| super::configured_default_host(doc));
198 let profile_name = host_runtime_config_from_doc(doc, &host)?.memory_profile;
199 Ok(ProfileSelection {
200 host: Some(host),
201 profile_name,
202 })
203}
204
205fn resolve_model_target(
206 target: &str,
207 reasoning_effort: Option<&str>,
208 executor: MemoryAiExecutor,
209) -> Result<ModelTarget> {
210 let target = target.trim();
211 if target.is_empty() {
212 bail!("model or preset must not be empty");
213 }
214 let lower = target.to_ascii_lowercase();
215 if let Some(preset) = MODEL_PRESETS.iter().find(|preset| preset.name == lower) {
216 if executor != MemoryAiExecutor::CodexCli {
217 bail!(
218 "model preset '{}' is for codex-cli profiles; pass an explicit model for {}",
219 preset.name,
220 executor.as_str()
221 );
222 }
223 if reasoning_effort.is_some() && preset.name == "auto" {
224 bail!("--reasoning cannot be used with `auto`");
225 }
226 return Ok(ModelTarget {
227 model: preset.model.to_string(),
228 reasoning_effort: reasoning_effort
229 .map(normalize_reasoning_effort)
230 .transpose()?
231 .or_else(|| preset.reasoning_effort.map(str::to_string)),
232 update_reasoning: true,
233 });
234 }
235
236 if lower == "auto" && executor != MemoryAiExecutor::CodexCli {
237 bail!("model `auto` is only supported for codex-cli profiles");
238 }
239 if lower == "auto" && reasoning_effort.is_some() {
240 bail!("--reasoning cannot be used with `auto`");
241 }
242 Ok(ModelTarget {
243 model: canonical_model_name(target),
244 reasoning_effort: reasoning_effort
245 .map(normalize_reasoning_effort)
246 .transpose()?,
247 update_reasoning: reasoning_effort.is_some() || lower == "auto",
248 })
249}
250
251fn canonical_model_name(model: &str) -> String {
252 match model.trim().to_ascii_lowercase().as_str() {
253 "5.4-mini" | "gpt5-4.mini" | "gpt-5-4-mini" => "gpt-5.4-mini".to_string(),
254 "5.2" | "gpt5.2" | "gpt-5-2" => "gpt-5.2".to_string(),
255 other => other.to_string(),
256 }
257}
258
259fn normalize_reasoning_effort(reasoning_effort: &str) -> Result<String> {
260 match reasoning_effort.trim().to_ascii_lowercase().as_str() {
261 "low" | "medium" | "high" => Ok(reasoning_effort.trim().to_ascii_lowercase()),
262 other => bail!("unknown reasoning effort '{other}'; expected low, medium, or high"),
263 }
264}
265
266fn profile_table_mut<'a>(
267 doc: &'a mut toml_edit::DocumentMut,
268 profile_name: &str,
269) -> Result<&'a mut Table> {
270 doc.get_mut("memory_ai")
271 .and_then(Item::as_table_mut)
272 .and_then(|table| table.get_mut("profiles"))
273 .and_then(Item::as_table_mut)
274 .and_then(|profiles| profiles.get_mut(profile_name))
275 .and_then(Item::as_table_mut)
276 .with_context(|| format!("missing [memory_ai.profiles.{profile_name}]"))
277}
278
279fn backup_path_for_config(path: &Path) -> PathBuf {
280 let mut backup = path.to_path_buf();
281 backup.set_extension("toml.bak");
282 backup
283}
284
285#[cfg(test)]
286mod tests {
287 use super::*;
288
289 fn with_config_path<T>(path: &Path, f: impl FnOnce() -> T) -> T {
290 let _guard = super::super::TEST_ENV_LOCK
291 .lock()
292 .expect("env lock should acquire");
293 let old = std::env::var("REMEM_CONFIG").ok();
294 unsafe { std::env::set_var("REMEM_CONFIG", path) };
295 let result = f();
296 match old {
297 Some(value) => unsafe { std::env::set_var("REMEM_CONFIG", value) },
298 None => unsafe { std::env::remove_var("REMEM_CONFIG") },
299 }
300 result
301 }
302
303 fn temp_config_path(label: &str) -> PathBuf {
304 std::env::temp_dir().join(format!(
305 "remem-model-{label}-{}-{}.toml",
306 std::process::id(),
307 chrono::Utc::now().timestamp_nanos_opt().unwrap_or_default()
308 ))
309 }
310
311 #[test]
312 fn model_use_preset_updates_codex_profile_and_backup() {
313 let path = temp_config_path("preset");
314 with_config_path(&path, || {
315 super::super::init_config().unwrap();
316 let change = set_model(Some(CODEX_HOST), None, "balanced", None, false).unwrap();
317 assert_eq!(change.old_model, DEFAULT_CODEX_MODEL);
318 assert_eq!(change.new_model, "gpt-5.4-mini");
319 assert_eq!(change.new_reasoning_effort.as_deref(), Some("medium"));
320 assert!(change.backup_path.as_ref().unwrap().exists());
321
322 let status = model_status(Some(CODEX_HOST), None).unwrap();
323 assert_eq!(status.reasoning_effort.as_deref(), Some("medium"));
324 });
325 let _ = std::fs::remove_file(&path);
326 let _ = std::fs::remove_file(backup_path_for_config(&path));
327 }
328
329 #[test]
330 fn model_use_dry_run_does_not_write() {
331 let path = temp_config_path("dry-run");
332 with_config_path(&path, || {
333 super::super::init_config().unwrap();
334 let change = set_model(Some(CODEX_HOST), None, "quality", None, true).unwrap();
335 assert!(change.dry_run);
336 assert_eq!(change.new_model, "gpt-5.2");
337
338 let status = model_status(Some(CODEX_HOST), None).unwrap();
339 assert_eq!(status.model, DEFAULT_CODEX_MODEL);
340 });
341 let _ = std::fs::remove_file(&path);
342 }
343
344 #[test]
345 fn model_rollback_restores_backup() {
346 let path = temp_config_path("rollback");
347 with_config_path(&path, || {
348 super::super::init_config().unwrap();
349 set_model(Some(CODEX_HOST), None, "quality", None, false).unwrap();
350 rollback_model_config().unwrap();
351
352 let status = model_status(Some(CODEX_HOST), None).unwrap();
353 assert_eq!(status.model, DEFAULT_CODEX_MODEL);
354 assert_eq!(status.reasoning_effort.as_deref(), None);
355 });
356 let _ = std::fs::remove_file(&path);
357 let _ = std::fs::remove_file(backup_path_for_config(&path));
358 }
359
360 #[test]
361 fn rollback_failure_preserves_active_config() -> Result<()> {
362 let path = temp_config_path("rollback-atomic-fail");
363 with_config_path(&path, || -> Result<()> {
364 let _atomic_guard = crate::atomic_file::failpoint_test_lock();
365 super::super::init_config()?;
366 set_model(Some(CODEX_HOST), None, "balanced", None, false)?;
367 let before = std::fs::read_to_string(&path)?;
368 crate::atomic_file::fail_next_rename_for_path_for_test(&path);
369
370 let err = rollback_model_config().expect_err("injected failure must abort rollback");
371 assert!(format!("{err:?}").contains("injected atomic write failure"));
372 assert_eq!(std::fs::read_to_string(&path)?, before);
373 crate::atomic_file::clear_failpoints_for_test();
374 Ok(())
375 })?;
376 let _ = std::fs::remove_file(&path);
377 let _ = std::fs::remove_file(backup_path_for_config(&path));
378 Ok(())
379 }
380
381 #[test]
382 fn mini_model_aliases_do_not_follow_default_model() {
383 assert_eq!(canonical_model_name("5.4-mini"), "gpt-5.4-mini");
384 assert_eq!(canonical_model_name("gpt-5-4-mini"), "gpt-5.4-mini");
385 assert_eq!(canonical_model_name("5.2"), DEFAULT_CODEX_MODEL);
386 }
387}