Skip to main content

lean_ctx/core/
data_dir.rs

1use std::path::PathBuf;
2
3const DATA_MARKERS: &[&str] = &["stats.json", "config.toml", "sessions"];
4
5/// Resolve the lean-ctx data directory.
6///
7/// Priority order (backward-compatible XDG migration):
8/// 1. `LEAN_CTX_DATA_DIR` env var (explicit override)
9/// 2. `~/.lean-ctx` if it has actual data (stats.json/config.toml/sessions)
10/// 3. `$XDG_CONFIG_HOME/lean-ctx` (XDG compliant, default `~/.config/lean-ctx`)
11///
12/// An empty `~/.lean-ctx/` directory does NOT trigger legacy mode — this prevents
13/// data directory splits when setup creates the dir before MCP writes stats.
14pub fn lean_ctx_data_dir() -> Result<PathBuf, String> {
15    if let Ok(dir) = std::env::var("LEAN_CTX_DATA_DIR") {
16        let trimmed = dir.trim();
17        if !trimmed.is_empty() {
18            let p = PathBuf::from(trimmed);
19            ensure_dir_permissions(&p);
20            return Ok(p);
21        }
22    }
23
24    let home = dirs::home_dir().ok_or_else(|| "Cannot determine home directory".to_string())?;
25
26    let legacy = home.join(".lean-ctx");
27    if legacy.exists() && has_data_files(&legacy) {
28        ensure_dir_permissions(&legacy);
29        return Ok(legacy);
30    }
31
32    let xdg_config = std::env::var("XDG_CONFIG_HOME")
33        .ok()
34        .filter(|s| !s.trim().is_empty())
35        .map_or_else(|| home.join(".config"), PathBuf::from);
36
37    let xdg_dir = xdg_config.join("lean-ctx");
38
39    if xdg_dir.exists() && has_data_files(&xdg_dir) {
40        ensure_dir_permissions(&xdg_dir);
41        return Ok(xdg_dir);
42    }
43
44    if legacy.exists() {
45        ensure_dir_permissions(&legacy);
46        return Ok(legacy);
47    }
48
49    ensure_dir_permissions(&xdg_dir);
50    Ok(xdg_dir)
51}
52
53fn has_data_files(dir: &std::path::Path) -> bool {
54    DATA_MARKERS.iter().any(|f| dir.join(f).exists())
55}
56
57/// Returns all known data directories that contain stats data.
58/// Used for migration and doctor diagnostics.
59pub fn all_data_dirs_with_stats() -> Vec<PathBuf> {
60    let mut dirs = Vec::new();
61    if let Some(home) = dirs::home_dir() {
62        let legacy = home.join(".lean-ctx");
63        if legacy.join("stats.json").exists() {
64            dirs.push(legacy);
65        }
66        let xdg = std::env::var("XDG_CONFIG_HOME")
67            .ok()
68            .filter(|s| !s.trim().is_empty())
69            .map_or_else(|| home.join(".config"), PathBuf::from)
70            .join("lean-ctx");
71        if xdg.join("stats.json").exists() && !dirs.contains(&xdg) {
72            dirs.push(xdg);
73        }
74    }
75    dirs
76}
77
78/// Detect and repair a data directory split.
79/// Returns the number of tokens migrated, or None if no split detected.
80pub fn migrate_if_split() -> Option<u64> {
81    let dirs = all_data_dirs_with_stats();
82    if dirs.len() < 2 {
83        return None;
84    }
85
86    let primary = lean_ctx_data_dir().ok()?;
87    let secondary = dirs.iter().find(|d| **d != primary)?;
88
89    let sec_content = std::fs::read_to_string(secondary.join("stats.json")).ok()?;
90    let sec_store: serde_json::Value = serde_json::from_str(&sec_content).ok()?;
91    let sec_commands = sec_store["total_commands"].as_u64().unwrap_or(0);
92    if sec_commands == 0 {
93        return None;
94    }
95
96    let primary_path = primary.join("stats.json");
97    if !primary_path.exists() {
98        let _ = std::fs::create_dir_all(&primary);
99        let _ = std::fs::copy(secondary.join("stats.json"), &primary_path);
100        let _ = std::fs::remove_file(secondary.join("stats.json"));
101        let tokens = sec_store["total_input_tokens"]
102            .as_u64()
103            .unwrap_or(0)
104            .saturating_sub(sec_store["total_output_tokens"].as_u64().unwrap_or(0));
105        return Some(tokens);
106    }
107
108    None
109}
110
111#[cfg(unix)]
112fn ensure_dir_permissions(path: &std::path::Path) {
113    use std::os::unix::fs::PermissionsExt;
114    if path.is_dir() {
115        let _ = std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700));
116    }
117}
118
119#[cfg(not(unix))]
120fn ensure_dir_permissions(_path: &std::path::Path) {}
121
122pub fn test_env_lock() -> std::sync::MutexGuard<'static, ()> {
123    use std::sync::{Mutex, OnceLock};
124    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
125    let mutex = LOCK.get_or_init(|| Mutex::new(()));
126    mutex
127        .lock()
128        .unwrap_or_else(std::sync::PoisonError::into_inner)
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134
135    #[test]
136    fn has_data_files_empty_dir() {
137        let dir = std::env::temp_dir().join("test_data_dir_empty");
138        let _ = std::fs::remove_dir_all(&dir);
139        let _ = std::fs::create_dir_all(&dir);
140        assert!(!has_data_files(&dir));
141        let _ = std::fs::remove_dir_all(&dir);
142    }
143
144    #[test]
145    fn has_data_files_with_stats() {
146        let dir = std::env::temp_dir().join("test_data_dir_stats");
147        let _ = std::fs::remove_dir_all(&dir);
148        let _ = std::fs::create_dir_all(&dir);
149        std::fs::write(dir.join("stats.json"), "{}").unwrap();
150        assert!(has_data_files(&dir));
151        let _ = std::fs::remove_dir_all(&dir);
152    }
153
154    #[test]
155    fn has_data_files_with_config() {
156        let dir = std::env::temp_dir().join("test_data_dir_config");
157        let _ = std::fs::remove_dir_all(&dir);
158        let _ = std::fs::create_dir_all(&dir);
159        std::fs::write(dir.join("config.toml"), "").unwrap();
160        assert!(has_data_files(&dir));
161        let _ = std::fs::remove_dir_all(&dir);
162    }
163
164    #[test]
165    fn has_data_files_with_sessions() {
166        let dir = std::env::temp_dir().join("test_data_dir_sessions");
167        let _ = std::fs::remove_dir_all(&dir);
168        let _ = std::fs::create_dir_all(&dir);
169        let _ = std::fs::create_dir_all(dir.join("sessions"));
170        assert!(has_data_files(&dir));
171        let _ = std::fs::remove_dir_all(&dir);
172    }
173
174    #[test]
175    fn lean_ctx_data_dir_env_override() {
176        let _lock = test_env_lock();
177        let dir = std::env::temp_dir().join("test_data_dir_env");
178        let _ = std::fs::remove_dir_all(&dir);
179        let _ = std::fs::create_dir_all(&dir);
180        std::env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
181        let result = lean_ctx_data_dir().unwrap();
182        assert_eq!(result, dir);
183        std::env::remove_var("LEAN_CTX_DATA_DIR");
184        let _ = std::fs::remove_dir_all(&dir);
185    }
186
187    #[test]
188    fn has_data_files_is_false_for_empty_dir() {
189        let dir = std::env::temp_dir().join("test_data_dir_no_data");
190        let _ = std::fs::remove_dir_all(&dir);
191        let _ = std::fs::create_dir_all(&dir);
192        std::fs::write(dir.join("random.txt"), "not a marker").unwrap();
193        assert!(!has_data_files(&dir));
194        let _ = std::fs::remove_dir_all(&dir);
195    }
196
197    #[test]
198    fn xdg_override_with_data_wins() {
199        let _lock = test_env_lock();
200
201        let xdg_base = std::env::temp_dir().join("test_xdg_override_wins");
202        let _ = std::fs::remove_dir_all(&xdg_base);
203        let xdg_dir = xdg_base.join("lean-ctx");
204        let _ = std::fs::create_dir_all(&xdg_dir);
205        std::fs::write(xdg_dir.join("stats.json"), r#"{"total_commands":1}"#).unwrap();
206
207        std::env::set_var("LEAN_CTX_DATA_DIR", "");
208        std::env::set_var("XDG_CONFIG_HOME", xdg_base.to_str().unwrap());
209
210        let result = lean_ctx_data_dir().unwrap();
211
212        std::env::remove_var("LEAN_CTX_DATA_DIR");
213        std::env::remove_var("XDG_CONFIG_HOME");
214
215        let home = dirs::home_dir().unwrap();
216        let legacy = home.join(".lean-ctx");
217        if !has_data_files(&legacy) {
218            assert_eq!(
219                result, xdg_dir,
220                "XDG with data should win when legacy has no data"
221            );
222        }
223
224        let _ = std::fs::remove_dir_all(&xdg_base);
225    }
226}