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    // Test sandbox (GL #512): without this, any unit test that triggers a
25    // store write (stats, savings ledger, context ledger, heatmap, ...)
26    // silently pollutes the developer's real ~/.lean-ctx — bounce events from
27    // test fixtures showed up as "today 61%" on the user dashboard. Tests that
28    // set LEAN_CTX_DATA_DIR keep full control (handled above); everyone else
29    // lands in a per-process temp dir and physically cannot touch real data.
30    #[cfg(test)]
31    {
32        static TEST_SANDBOX: std::sync::OnceLock<PathBuf> = std::sync::OnceLock::new();
33        let dir = TEST_SANDBOX.get_or_init(|| {
34            let d = std::env::temp_dir().join(format!("lean-ctx-testdata-{}", std::process::id()));
35            let _ = std::fs::create_dir_all(&d);
36            d
37        });
38        Ok(dir.clone())
39    }
40
41    #[cfg(not(test))]
42    {
43        resolve_home_data_dir()
44    }
45}
46
47/// Home-based resolution (legacy `~/.lean-ctx` vs XDG). Split out so the
48/// priority rules stay unit-testable despite the test sandbox above.
49fn resolve_home_data_dir() -> Result<PathBuf, String> {
50    let home = dirs::home_dir().ok_or_else(|| "Cannot determine home directory".to_string())?;
51
52    let legacy = home.join(".lean-ctx");
53    if legacy.exists() && has_data_files(&legacy) {
54        ensure_dir_permissions(&legacy);
55        return Ok(legacy);
56    }
57
58    let xdg_config = std::env::var("XDG_CONFIG_HOME")
59        .ok()
60        .filter(|s| !s.trim().is_empty())
61        .map_or_else(|| home.join(".config"), PathBuf::from);
62
63    let xdg_dir = xdg_config.join("lean-ctx");
64
65    if xdg_dir.exists() && has_data_files(&xdg_dir) {
66        ensure_dir_permissions(&xdg_dir);
67        return Ok(xdg_dir);
68    }
69
70    if legacy.exists() {
71        ensure_dir_permissions(&legacy);
72        return Ok(legacy);
73    }
74
75    ensure_dir_permissions(&xdg_dir);
76    Ok(xdg_dir)
77}
78
79fn has_data_files(dir: &std::path::Path) -> bool {
80    DATA_MARKERS.iter().any(|f| dir.join(f).exists())
81}
82
83/// Returns all known data directories that contain stats data.
84/// Used for migration and doctor diagnostics.
85pub fn all_data_dirs_with_stats() -> Vec<PathBuf> {
86    let mut dirs = Vec::new();
87    if let Some(home) = dirs::home_dir() {
88        let legacy = home.join(".lean-ctx");
89        if legacy.join("stats.json").exists() {
90            dirs.push(legacy);
91        }
92        let xdg = std::env::var("XDG_CONFIG_HOME")
93            .ok()
94            .filter(|s| !s.trim().is_empty())
95            .map_or_else(|| home.join(".config"), PathBuf::from)
96            .join("lean-ctx");
97        if xdg.join("stats.json").exists() && !dirs.contains(&xdg) {
98            dirs.push(xdg);
99        }
100    }
101    dirs
102}
103
104/// Detect and repair a data directory split.
105/// Returns the number of tokens migrated, or None if no split detected.
106pub fn migrate_if_split() -> Option<u64> {
107    let dirs = all_data_dirs_with_stats();
108    if dirs.len() < 2 {
109        return None;
110    }
111
112    let primary = lean_ctx_data_dir().ok()?;
113    let secondary = dirs.iter().find(|d| **d != primary)?;
114
115    let sec_content = std::fs::read_to_string(secondary.join("stats.json")).ok()?;
116    let sec_store: serde_json::Value = serde_json::from_str(&sec_content).ok()?;
117    let sec_commands = sec_store["total_commands"].as_u64().unwrap_or(0);
118    if sec_commands == 0 {
119        return None;
120    }
121
122    let primary_path = primary.join("stats.json");
123    if !primary_path.exists() {
124        let _ = std::fs::create_dir_all(&primary);
125        let _ = std::fs::copy(secondary.join("stats.json"), &primary_path);
126        let _ = std::fs::remove_file(secondary.join("stats.json"));
127        let tokens = sec_store["total_input_tokens"]
128            .as_u64()
129            .unwrap_or(0)
130            .saturating_sub(sec_store["total_output_tokens"].as_u64().unwrap_or(0));
131        return Some(tokens);
132    }
133
134    None
135}
136
137#[cfg(unix)]
138fn ensure_dir_permissions(path: &std::path::Path) {
139    use std::os::unix::fs::PermissionsExt;
140    if path.is_dir() {
141        let _ = std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700));
142    }
143}
144
145#[cfg(not(unix))]
146fn ensure_dir_permissions(_path: &std::path::Path) {}
147
148pub fn test_env_lock() -> std::sync::MutexGuard<'static, ()> {
149    use std::sync::{Mutex, OnceLock};
150    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
151    let mutex = LOCK.get_or_init(|| Mutex::new(()));
152    mutex
153        .lock()
154        .unwrap_or_else(std::sync::PoisonError::into_inner)
155}
156
157/// RAII data-dir isolation for tests (GL #556): holds `test_env_lock` for
158/// the guard's lifetime, points `LEAN_CTX_DATA_DIR` at a fresh temp dir and
159/// restores the env on drop — even on panic, so a failing test cannot leak
160/// the override into others. Use this instead of hand-rolled
161/// `set_var`/`remove_var` pairs whenever a test needs an empty, private
162/// data dir (the shared per-process sandbox is NOT empty: parallel tests
163/// write stores like feedback, bandit and sessions into it).
164#[cfg(test)]
165pub struct IsolatedDataDir {
166    tmp: tempfile::TempDir,
167    _guard: std::sync::MutexGuard<'static, ()>,
168}
169
170#[cfg(test)]
171impl IsolatedDataDir {
172    pub fn path(&self) -> &std::path::Path {
173        self.tmp.path()
174    }
175}
176
177#[cfg(test)]
178impl Drop for IsolatedDataDir {
179    fn drop(&mut self) {
180        // Struct Drop runs before field drops, so the env is restored while
181        // the lock is still held.
182        std::env::remove_var("LEAN_CTX_DATA_DIR");
183    }
184}
185
186#[cfg(test)]
187pub fn isolated_data_dir() -> IsolatedDataDir {
188    let guard = test_env_lock();
189    let tmp = tempfile::tempdir().expect("tempdir for isolated data dir");
190    std::env::set_var("LEAN_CTX_DATA_DIR", tmp.path());
191    IsolatedDataDir { tmp, _guard: guard }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197
198    #[test]
199    fn has_data_files_empty_dir() {
200        let dir = std::env::temp_dir().join("test_data_dir_empty");
201        let _ = std::fs::remove_dir_all(&dir);
202        let _ = std::fs::create_dir_all(&dir);
203        assert!(!has_data_files(&dir));
204        let _ = std::fs::remove_dir_all(&dir);
205    }
206
207    #[test]
208    fn has_data_files_with_stats() {
209        let dir = std::env::temp_dir().join("test_data_dir_stats");
210        let _ = std::fs::remove_dir_all(&dir);
211        let _ = std::fs::create_dir_all(&dir);
212        std::fs::write(dir.join("stats.json"), "{}").unwrap();
213        assert!(has_data_files(&dir));
214        let _ = std::fs::remove_dir_all(&dir);
215    }
216
217    #[test]
218    fn has_data_files_with_config() {
219        let dir = std::env::temp_dir().join("test_data_dir_config");
220        let _ = std::fs::remove_dir_all(&dir);
221        let _ = std::fs::create_dir_all(&dir);
222        std::fs::write(dir.join("config.toml"), "").unwrap();
223        assert!(has_data_files(&dir));
224        let _ = std::fs::remove_dir_all(&dir);
225    }
226
227    #[test]
228    fn has_data_files_with_sessions() {
229        let dir = std::env::temp_dir().join("test_data_dir_sessions");
230        let _ = std::fs::remove_dir_all(&dir);
231        let _ = std::fs::create_dir_all(&dir);
232        let _ = std::fs::create_dir_all(dir.join("sessions"));
233        assert!(has_data_files(&dir));
234        let _ = std::fs::remove_dir_all(&dir);
235    }
236
237    #[test]
238    fn lean_ctx_data_dir_env_override() {
239        let _lock = test_env_lock();
240        let dir = std::env::temp_dir().join("test_data_dir_env");
241        let _ = std::fs::remove_dir_all(&dir);
242        let _ = std::fs::create_dir_all(&dir);
243        std::env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
244        let result = lean_ctx_data_dir().unwrap();
245        assert_eq!(result, dir);
246        std::env::remove_var("LEAN_CTX_DATA_DIR");
247        let _ = std::fs::remove_dir_all(&dir);
248    }
249
250    #[test]
251    fn has_data_files_is_false_for_empty_dir() {
252        let dir = std::env::temp_dir().join("test_data_dir_no_data");
253        let _ = std::fs::remove_dir_all(&dir);
254        let _ = std::fs::create_dir_all(&dir);
255        std::fs::write(dir.join("random.txt"), "not a marker").unwrap();
256        assert!(!has_data_files(&dir));
257        let _ = std::fs::remove_dir_all(&dir);
258    }
259
260    #[test]
261    fn xdg_override_with_data_wins() {
262        let _lock = test_env_lock();
263
264        let xdg_base = std::env::temp_dir().join("test_xdg_override_wins");
265        let _ = std::fs::remove_dir_all(&xdg_base);
266        let xdg_dir = xdg_base.join("lean-ctx");
267        let _ = std::fs::create_dir_all(&xdg_dir);
268        std::fs::write(xdg_dir.join("stats.json"), r#"{"total_commands":1}"#).unwrap();
269
270        std::env::set_var("LEAN_CTX_DATA_DIR", "");
271        std::env::set_var("XDG_CONFIG_HOME", xdg_base.to_str().unwrap());
272
273        // Calls the home resolver directly: lean_ctx_data_dir() is sandboxed
274        // under cfg(test) (GL #512) and would short-circuit before XDG logic.
275        let result = resolve_home_data_dir().unwrap();
276
277        std::env::remove_var("LEAN_CTX_DATA_DIR");
278        std::env::remove_var("XDG_CONFIG_HOME");
279
280        let home = dirs::home_dir().unwrap();
281        let legacy = home.join(".lean-ctx");
282        if !has_data_files(&legacy) {
283            assert_eq!(
284                result, xdg_dir,
285                "XDG with data should win when legacy has no data"
286            );
287        }
288
289        let _ = std::fs::remove_dir_all(&xdg_base);
290    }
291}