Skip to main content

codex_helper_core/
local_log_store.rs

1use std::fs::{self, File, OpenOptions};
2use std::io::{self, Write};
3use std::path::{Path, PathBuf};
4use std::time::{SystemTime, UNIX_EPOCH};
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub struct LogRetention {
8    pub max_bytes: u64,
9    pub max_files: usize,
10}
11
12impl LogRetention {
13    pub const fn new(max_bytes: u64, max_files: usize) -> Self {
14        Self {
15            max_bytes,
16            max_files,
17        }
18    }
19
20    pub fn from_env(
21        max_bytes_key: &str,
22        max_files_key: &str,
23        default_max_bytes: u64,
24        default_max_files: usize,
25    ) -> Self {
26        Self {
27            max_bytes: parse_u64_env(max_bytes_key).unwrap_or(default_max_bytes),
28            max_files: parse_usize_env(max_files_key).unwrap_or(default_max_files),
29        }
30    }
31
32    pub fn enabled(self) -> bool {
33        self.max_bytes > 0 && self.max_files > 0
34    }
35
36    fn rotated_budget_bytes(self) -> u64 {
37        self.max_bytes.saturating_mul(self.max_files as u64)
38    }
39}
40
41#[derive(Debug)]
42pub struct RotatedLogFile {
43    pub path: PathBuf,
44    pub bytes: u64,
45    pub modified: SystemTime,
46}
47
48pub fn repair_log(path: impl AsRef<Path>, retention: LogRetention) {
49    let path = path.as_ref();
50    rotate_and_prune_if_needed(path, retention);
51    prune_rotated_logs(path, retention);
52}
53
54pub fn append_line(path: impl AsRef<Path>, retention: LogRetention, line: &str) -> io::Result<()> {
55    let path = path.as_ref();
56    if let Some(parent) = path.parent() {
57        fs::create_dir_all(parent)?;
58    }
59    rotate_and_prune_if_needed(path, retention);
60    let mut file = OpenOptions::new().create(true).append(true).open(path)?;
61    writeln!(file, "{line}")?;
62    Ok(())
63}
64
65pub fn rotate_and_prune_if_needed(path: impl AsRef<Path>, retention: LogRetention) {
66    let path = path.as_ref();
67    if !retention.enabled() {
68        return;
69    }
70    let Ok(meta) = fs::metadata(path) else {
71        return;
72    };
73    if meta.len() < retention.max_bytes {
74        prune_rotated_logs(path, retention);
75        return;
76    }
77    if rotate_log_path(path).is_ok() {
78        prune_rotated_logs(path, retention);
79    }
80}
81
82pub fn prune_rotated_logs(path: impl AsRef<Path>, retention: LogRetention) {
83    let path = path.as_ref();
84    prune_rotated_logs_with_remover(path, retention, |candidate| fs::remove_file(candidate));
85}
86
87fn prune_rotated_logs_with_remover(
88    path: &Path,
89    retention: LogRetention,
90    mut remove_file: impl FnMut(&Path) -> io::Result<()>,
91) {
92    if !retention.enabled() {
93        return;
94    }
95    let mut rotated = collect_rotated_logs(path);
96    let mut total_bytes = rotated
97        .iter()
98        .fold(0_u64, |acc, file| acc.saturating_add(file.bytes));
99    let budget_bytes = retention.rotated_budget_bytes();
100
101    while !rotated.is_empty() && (rotated.len() > retention.max_files || total_bytes > budget_bytes)
102    {
103        let file = rotated.remove(0);
104        match remove_file(&file.path) {
105            Ok(()) => {
106                total_bytes = total_bytes.saturating_sub(file.bytes);
107            }
108            Err(err) if err.kind() == io::ErrorKind::NotFound => {
109                total_bytes = total_bytes.saturating_sub(file.bytes);
110            }
111            Err(_) => {}
112        }
113    }
114}
115
116pub fn collect_rotated_logs(path: impl AsRef<Path>) -> Vec<RotatedLogFile> {
117    let path = path.as_ref();
118    let Some(parent) = path.parent() else {
119        return Vec::new();
120    };
121    let Ok(rd) = fs::read_dir(parent) else {
122        return Vec::new();
123    };
124    let mut rotated: Vec<RotatedLogFile> = rd
125        .filter_map(|entry| entry.ok())
126        .filter_map(|entry| {
127            let entry_path = entry.path();
128            if !is_rotated_log_path(path, &entry_path) {
129                return None;
130            }
131            let meta = entry.metadata().ok()?;
132            Some(RotatedLogFile {
133                path: entry_path,
134                bytes: meta.len(),
135                modified: meta.modified().unwrap_or(SystemTime::UNIX_EPOCH),
136            })
137        })
138        .collect();
139    rotated.sort_by(|left, right| {
140        left.modified
141            .cmp(&right.modified)
142            .then_with(|| left.path.cmp(&right.path))
143    });
144    rotated
145}
146
147pub struct RotatingLogWriter {
148    path: PathBuf,
149    retention: LogRetention,
150    file: Option<File>,
151    current_len: u64,
152}
153
154impl RotatingLogWriter {
155    pub fn new(path: PathBuf, retention: LogRetention) -> Self {
156        let current_len = fs::metadata(&path).map(|meta| meta.len()).unwrap_or(0);
157        Self {
158            path,
159            retention,
160            file: None,
161            current_len,
162        }
163    }
164
165    fn ensure_file(&mut self) -> io::Result<&mut File> {
166        if self.file.is_none() {
167            if let Some(parent) = self.path.parent() {
168                fs::create_dir_all(parent)?;
169            }
170            let file = OpenOptions::new()
171                .create(true)
172                .append(true)
173                .open(&self.path)?;
174            self.current_len = file.metadata().map(|meta| meta.len()).unwrap_or(0);
175            self.file = Some(file);
176        }
177        self.file
178            .as_mut()
179            .ok_or_else(|| io::Error::other("bounded log file was not opened"))
180    }
181
182    fn rotate_before_write(&mut self, incoming_len: usize) {
183        if !self.retention.enabled() || self.current_len == 0 {
184            return;
185        }
186        let incoming_len = incoming_len as u64;
187        if self.current_len.saturating_add(incoming_len) < self.retention.max_bytes {
188            return;
189        }
190        self.file.take();
191        if rotate_log_path(&self.path).is_ok() {
192            self.current_len = 0;
193            prune_rotated_logs(&self.path, self.retention);
194        } else {
195            self.current_len = fs::metadata(&self.path).map(|meta| meta.len()).unwrap_or(0);
196        }
197    }
198}
199
200impl Write for RotatingLogWriter {
201    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
202        self.rotate_before_write(buf.len());
203        let file = self.ensure_file()?;
204        let written = file.write(buf)?;
205        self.current_len = self.current_len.saturating_add(written as u64);
206        Ok(written)
207    }
208
209    fn flush(&mut self) -> io::Result<()> {
210        if let Some(file) = self.file.as_mut() {
211            file.flush()
212        } else {
213            Ok(())
214        }
215    }
216}
217
218fn rotate_log_path(path: &Path) -> io::Result<()> {
219    for attempt in 0..100 {
220        let rotated_path = rotated_path_for(path, attempt);
221        if rotated_path.exists() {
222            continue;
223        }
224        return fs::rename(path, rotated_path);
225    }
226    Err(io::Error::new(
227        io::ErrorKind::AlreadyExists,
228        "could not allocate rotated log file name",
229    ))
230}
231
232fn rotated_path_for(path: &Path, attempt: u32) -> PathBuf {
233    let ts = SystemTime::now()
234        .duration_since(UNIX_EPOCH)
235        .map(|d| d.as_nanos())
236        .unwrap_or(0);
237    let attempt_suffix = if attempt == 0 {
238        String::new()
239    } else {
240        format!(".{attempt}")
241    };
242    let file_name = path.file_name().and_then(|s| s.to_str()).unwrap_or("log");
243
244    if path.extension().and_then(|s| s.to_str()) == Some("jsonl")
245        && let Some(stem) = path.file_stem().and_then(|s| s.to_str())
246    {
247        return path.with_file_name(format!("{stem}.{ts}{attempt_suffix}.jsonl"));
248    }
249
250    path.with_file_name(format!("{file_name}.{ts}{attempt_suffix}"))
251}
252
253fn is_rotated_log_path(active_path: &Path, candidate_path: &Path) -> bool {
254    if active_path == candidate_path {
255        return false;
256    }
257    let Some(name) = candidate_path.file_name().and_then(|s| s.to_str()) else {
258        return false;
259    };
260    let Some(active_name) = active_path.file_name().and_then(|s| s.to_str()) else {
261        return false;
262    };
263
264    if active_path.extension().and_then(|s| s.to_str()) == Some("jsonl")
265        && let Some(stem) = active_path.file_stem().and_then(|s| s.to_str())
266    {
267        return name.starts_with(&format!("{stem}.")) && name.ends_with(".jsonl");
268    }
269
270    name.starts_with(&format!("{active_name}."))
271}
272
273fn parse_u64_env(key: &str) -> Option<u64> {
274    std::env::var(key)
275        .ok()
276        .and_then(|s| s.trim().parse::<u64>().ok())
277        .filter(|&n| n > 0)
278}
279
280fn parse_usize_env(key: &str) -> Option<usize> {
281    std::env::var(key)
282        .ok()
283        .and_then(|s| s.trim().parse::<usize>().ok())
284        .filter(|&n| n > 0)
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290
291    fn temp_log_dir(test_name: &str) -> PathBuf {
292        let dir = std::env::temp_dir().join(format!(
293            "codex-helper-local-log-store-{test_name}-{}-{}",
294            std::process::id(),
295            SystemTime::now()
296                .duration_since(UNIX_EPOCH)
297                .map(|d| d.as_nanos())
298                .unwrap_or(0)
299        ));
300        let _ = fs::remove_dir_all(&dir);
301        fs::create_dir_all(&dir).expect("create temp log dir");
302        dir
303    }
304
305    fn write_test_file(path: &Path, bytes: usize) {
306        let mut file = File::create(path).expect("create test log file");
307        file.write_all(&vec![b'x'; bytes])
308            .expect("write test log file");
309        file.flush().expect("flush test log file");
310    }
311
312    #[test]
313    fn repair_log_removes_legacy_runtime_rotated_file_over_budget() {
314        let dir = temp_log_dir("legacy-runtime-rotated-budget");
315        let active = dir.join("runtime.log");
316        let legacy = dir.join("runtime.log.legacy");
317        write_test_file(&legacy, 64);
318
319        repair_log(&active, LogRetention::new(10, 2));
320
321        assert!(
322            !legacy.exists(),
323            "oversized legacy rotated runtime log should be removed"
324        );
325        let _ = fs::remove_dir_all(dir);
326    }
327
328    #[test]
329    fn repair_log_rotates_and_prunes_oversized_active_file() {
330        let dir = temp_log_dir("active-oversized");
331        let active = dir.join("runtime.log");
332        write_test_file(&active, 24);
333
334        repair_log(&active, LogRetention::new(10, 2));
335
336        assert!(
337            !active.exists(),
338            "oversized active log should be rotated away before tracing starts"
339        );
340        assert!(
341            collect_rotated_logs(&active).is_empty(),
342            "rotated oversized active log should be pruned by total budget"
343        );
344        let _ = fs::remove_dir_all(dir);
345    }
346
347    #[test]
348    fn prune_rotated_logs_keeps_budget_pressure_when_delete_fails() {
349        let dir = temp_log_dir("delete-fail-budget-pressure");
350        let active = dir.join("runtime.log");
351        let locked = dir.join("runtime.log.0-locked");
352        let spill = dir.join("runtime.log.1-spill");
353        write_test_file(&locked, 64);
354        std::thread::sleep(std::time::Duration::from_millis(10));
355        write_test_file(&spill, 8);
356
357        prune_rotated_logs_with_remover(
358            &active,
359            LogRetention::new(10, 1),
360            |path| -> io::Result<()> {
361                if path == locked {
362                    return Err(io::Error::new(io::ErrorKind::PermissionDenied, "locked"));
363                }
364                fs::remove_file(path)
365            },
366        );
367
368        assert!(
369            locked.exists(),
370            "failed delete should leave the locked rotated log for a later repair"
371        );
372        assert!(
373            !spill.exists(),
374            "failed delete must not be counted as recovered budget"
375        );
376
377        prune_rotated_logs(&active, LogRetention::new(10, 1));
378        assert!(
379            !locked.exists(),
380            "next repair should remove the previously locked oversized log"
381        );
382        let _ = fs::remove_dir_all(dir);
383    }
384
385    #[test]
386    fn rotating_log_writer_rotates_while_process_is_running() {
387        let dir = temp_log_dir("writer-rotates");
388        let active = dir.join("runtime.log");
389        let retention = LogRetention::new(16, 2);
390        let mut writer = RotatingLogWriter::new(active.clone(), retention);
391
392        writer
393            .write_all(b"first-line-000\n")
394            .expect("write first line");
395        writer
396            .write_all(b"second-line-00\n")
397            .expect("write second line");
398        writer.flush().expect("flush runtime log writer");
399        drop(writer);
400
401        assert!(active.exists(), "new active log should exist");
402        assert!(
403            fs::metadata(&active).expect("active metadata").len() <= retention.max_bytes,
404            "active log should stay under the configured size after rotation"
405        );
406        assert_eq!(
407            collect_rotated_logs(&active).len(),
408            1,
409            "first active log should be rotated during the same process"
410        );
411        let _ = fs::remove_dir_all(dir);
412    }
413
414    #[test]
415    fn append_line_rotates_jsonl_before_extension() {
416        let dir = temp_log_dir("jsonl-rotation");
417        let active = dir.join("control_trace.jsonl");
418        write_test_file(&active, 24);
419
420        append_line(&active, LogRetention::new(16, 2), "{\"ok\":true}").expect("append jsonl line");
421
422        assert!(active.exists(), "append should recreate active jsonl file");
423        let rotated = collect_rotated_logs(&active);
424        assert_eq!(rotated.len(), 1);
425        let rotated_name = rotated[0]
426            .path
427            .file_name()
428            .and_then(|name| name.to_str())
429            .expect("rotated file name");
430        assert!(rotated_name.starts_with("control_trace."));
431        assert!(rotated_name.ends_with(".jsonl"));
432        let _ = fs::remove_dir_all(dir);
433    }
434}