rusty_stack/
benchmark_logs.rs

1use std::env;
2use std::fs;
3use std::path::{Path, PathBuf};
4
5const BENCHMARK_RESULTS_MARKER: &str = "---BENCHMARK_RESULTS_START---";
6
7/// Returns benchmark log directories mirroring the scripts/lib/benchmark_common.sh logic.
8/// Precedence: MLSTACK_LOG_DIR (when usable), then $HOME/.rusty-stack/logs, then ${TMPDIR:-/tmp}/rusty-stack/logs.
9pub fn benchmark_log_directories() -> Vec<PathBuf> {
10    let mut dirs = Vec::new();
11
12    if let Ok(env_dir) = env::var("MLSTACK_LOG_DIR") {
13        let trimmed = env_dir.trim();
14        if !trimmed.is_empty() {
15            let path = PathBuf::from(trimmed);
16            if let Some(dir) = existing_log_dir(&path) {
17                push_unique(&mut dirs, dir);
18            }
19        }
20    }
21
22    for home in candidate_log_homes() {
23        let path = Path::new(&home).join(".rusty-stack").join("logs");
24        if let Some(dir) = existing_log_dir(&path) {
25            push_unique(&mut dirs, dir);
26        }
27        // Codex/Claude environments often write benchmark logs here.
28        let claude_tmp = Path::new(&home)
29            .join(".claude")
30            .join("tmp")
31            .join("rusty-stack")
32            .join("logs");
33        if let Some(dir) = existing_log_dir(&claude_tmp) {
34            push_unique(&mut dirs, dir);
35        }
36    }
37
38    let tmp_root = env::var("TMPDIR").unwrap_or_else(|_| "/tmp".to_string());
39    if !tmp_root.trim().is_empty() {
40        let path = Path::new(&tmp_root).join("rusty-stack").join("logs");
41        if let Some(dir) = existing_log_dir(&path) {
42            push_unique(&mut dirs, dir);
43        }
44    }
45
46    dirs
47}
48
49/// Collects log files that contain any of the provided patterns, sorted by modification time.
50pub fn collect_matching_logs(log_dirs: &[PathBuf], patterns: &[&str]) -> Vec<PathBuf> {
51    let mut matches = Vec::new();
52    for dir in log_dirs {
53        if let Ok(entries) = fs::read_dir(dir) {
54            for entry in entries.filter_map(|entry| entry.ok()) {
55                let name = entry.file_name().to_string_lossy().into_owned();
56                if patterns.iter().any(|pat| name.contains(pat)) {
57                    let path = entry.path();
58                    if path.is_file() {
59                        matches.push(path);
60                    }
61                }
62            }
63        }
64    }
65
66    matches.sort_by_key(file_modified);
67    matches
68}
69
70/// Finds the most recent log file across the provided directories that matches the pattern.
71/// If any `.json` file is available, prefer the most recent JSON result.
72pub fn find_latest_log_in_dirs(log_dirs: &[PathBuf], pattern: &str) -> Option<PathBuf> {
73    let matches = collect_matching_logs(log_dirs, &[pattern]);
74    matches
75        .iter()
76        .filter(|path| is_json_log(path))
77        .max_by_key(|path| file_modified(path))
78        .cloned()
79        .or_else(|| matches.last().cloned())
80}
81
82pub fn is_json_log(path: &Path) -> bool {
83    path.extension()
84        .and_then(|ext| ext.to_str())
85        .map(|ext| ext.eq_ignore_ascii_case("json"))
86        .unwrap_or(false)
87}
88
89/// Extract a benchmark JSON object from raw log content.
90/// Handles plain JSON files, marker-delimited output, and mixed logs.
91pub fn extract_benchmark_json_value(contents: &str) -> Option<serde_json::Value> {
92    let trimmed = contents.trim();
93    if trimmed.is_empty() {
94        return None;
95    }
96
97    if let Ok(value) = serde_json::from_str::<serde_json::Value>(trimmed) {
98        return Some(value);
99    }
100
101    // Prefer marker matches that appear at line-start to avoid marker-like text inside JSON strings.
102    let marker_positions: Vec<usize> = contents
103        .match_indices(BENCHMARK_RESULTS_MARKER)
104        .map(|(idx, _)| idx)
105        .filter(|idx| marker_is_line_start(contents, *idx))
106        .collect();
107
108    for marker_idx in marker_positions.into_iter().rev() {
109        let search = &contents[marker_idx + BENCHMARK_RESULTS_MARKER.len()..];
110        if let Some(value) = parse_best_json_object(search) {
111            return Some(value);
112        }
113    }
114
115    parse_best_json_object(contents)
116}
117
118/// Extract benchmark JSON as a normalized string.
119pub fn extract_benchmark_json_string(contents: &str) -> Option<String> {
120    extract_benchmark_json_value(contents).and_then(|value| serde_json::to_string(&value).ok())
121}
122
123fn marker_is_line_start(contents: &str, marker_idx: usize) -> bool {
124    if marker_idx == 0 {
125        return true;
126    }
127    contents[..marker_idx].ends_with('\n')
128}
129
130fn parse_best_json_object(input: &str) -> Option<serde_json::Value> {
131    let mut best: Option<(usize, usize, serde_json::Value)> = None;
132
133    for (idx, ch) in input.char_indices() {
134        if ch != '{' {
135            continue;
136        }
137        let Some(end) = find_balanced_json_object_end(input, idx) else {
138            continue;
139        };
140
141        let candidate = &input[idx..end];
142        let Ok(value) = serde_json::from_str::<serde_json::Value>(candidate) else {
143            continue;
144        };
145
146        let score = benchmark_json_score(&value);
147        let candidate_len = end.saturating_sub(idx);
148        let is_better = best
149            .as_ref()
150            .map(|(best_score, best_len, _)| {
151                score > *best_score || (score == *best_score && candidate_len > *best_len)
152            })
153            .unwrap_or(true);
154        if is_better {
155            best = Some((score, candidate_len, value));
156        }
157    }
158
159    best.map(|(_, _, value)| value)
160}
161
162fn benchmark_json_score(value: &serde_json::Value) -> usize {
163    let mut score = 0usize;
164    if let Some(obj) = value.as_object() {
165        if obj.get("results").and_then(|v| v.as_object()).is_some() {
166            score += 200;
167        }
168        if obj.get("metrics").and_then(|v| v.as_object()).is_some() {
169            score += 100;
170        }
171        if obj.get("errors").and_then(|v| v.as_array()).is_some() {
172            score += 20;
173        }
174        if obj.get("success").and_then(|v| v.as_bool()).is_some() {
175            score += 15;
176        }
177        if obj.get("name").and_then(|v| v.as_str()).is_some() {
178            score += 10;
179        }
180        score += obj.len();
181    }
182
183    score
184}
185
186fn find_balanced_json_object_end(input: &str, start: usize) -> Option<usize> {
187    let mut depth = 0usize;
188    let mut in_string = false;
189    let mut escaped = false;
190
191    for (offset, ch) in input[start..].char_indices() {
192        if in_string {
193            if escaped {
194                escaped = false;
195                continue;
196            }
197            match ch {
198                '\\' => escaped = true,
199                '"' => in_string = false,
200                _ => {}
201            }
202            continue;
203        }
204
205        match ch {
206            '"' => in_string = true,
207            '{' => depth += 1,
208            '}' => {
209                if depth == 0 {
210                    return None;
211                }
212                depth -= 1;
213                if depth == 0 {
214                    return Some(start + offset + 1);
215                }
216            }
217            _ => {}
218        }
219    }
220
221    None
222}
223
224fn existing_log_dir(path: &Path) -> Option<PathBuf> {
225    if path.is_dir() {
226        Some(path.to_path_buf())
227    } else {
228        None
229    }
230}
231
232fn user_home_from_passwd(user_name: &str) -> Option<String> {
233    let target = user_name.trim();
234    if target.is_empty() {
235        return None;
236    }
237    let passwd = fs::read_to_string("/etc/passwd").ok()?;
238    passwd.lines().find_map(|line| {
239        if line.trim().is_empty() || line.starts_with('#') {
240            return None;
241        }
242        let fields: Vec<&str> = line.split(':').collect();
243        if fields.len() < 6 || fields[0] != target {
244            return None;
245        }
246        let home = fields[5].trim();
247        if home.is_empty() {
248            None
249        } else {
250            Some(home.to_string())
251        }
252    })
253}
254
255fn candidate_log_homes() -> Vec<String> {
256    let mut homes = Vec::new();
257
258    if let Ok(home) = env::var("MLSTACK_USER_HOME") {
259        let trimmed = home.trim();
260        if !trimmed.is_empty() {
261            homes.push(trimmed.to_string());
262        }
263    }
264    if let Ok(home) = env::var("HOME") {
265        let trimmed = home.trim();
266        if !trimmed.is_empty() {
267            homes.push(trimmed.to_string());
268        }
269    }
270    for key in ["SUDO_USER", "USER", "LOGNAME"] {
271        if let Ok(user_name) = env::var(key) {
272            if let Some(home) = user_home_from_passwd(&user_name) {
273                homes.push(home);
274            }
275        }
276    }
277
278    if let Ok(passwd) = fs::read_to_string("/etc/passwd") {
279        for line in passwd.lines() {
280            if line.trim().is_empty() || line.starts_with('#') {
281                continue;
282            }
283            let fields: Vec<&str> = line.split(':').collect();
284            if fields.len() < 6 {
285                continue;
286            }
287            let home = fields[5].trim();
288            if home.is_empty() {
289                continue;
290            }
291            let home_path = Path::new(home);
292            if home_path.join(".mlstack_env").is_file() || home_path.join(".mlstack").is_dir() {
293                homes.push(home.to_string());
294            }
295        }
296    }
297
298    homes.sort();
299    homes.dedup();
300    homes
301}
302
303fn push_unique(dirs: &mut Vec<PathBuf>, candidate: PathBuf) {
304    if !dirs.iter().any(|existing| existing == &candidate) {
305        dirs.push(candidate);
306    }
307}
308
309fn file_modified(path: &PathBuf) -> Option<std::time::SystemTime> {
310    fs::metadata(path)
311        .ok()
312        .and_then(|meta| meta.modified().ok())
313}
rusty_stack/benchmark_logs.rs

rusty_stack/
benchmark_logs.rs
