Skip to main content

borg_trace/
omp.rs

1//! omp (oh-my-pi) sessions: `~/.omp/agent/sessions/<slug>/<ts>_<id>.jsonl`.
2
3use std::collections::HashMap;
4use std::io::{BufRead, BufReader};
5use std::path::Path;
6use std::sync::LazyLock;
7
8use regex::Regex;
9use serde_json::Value;
10
11use crate::{Event, Harness, Kind, Session, approx_tokens, parse_ts};
12
13/// File headers inside omp's patch format: `[src/x.rs#06FE]` or `*** Update File: src/x.rs`.
14static PATCH_FILE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"(?m)^(?:\[([^\]#\n]+)(?:#[^\]\n]*)?\]|\*\*\* (?:Update|Add) File: (.+))$").unwrap());
15
16fn local_path(cwd: &str, p: &str) -> Option<String> {
17    let p = p.strip_prefix("local://").unwrap_or(p);
18    if p.contains("://") {
19        return None;
20    }
21    // `read` accepts `path:10-40` line selectors.
22    let p = match p.rsplit_once(':') {
23        Some((head, tail)) if tail.chars().next().is_some_and(|c| c.is_ascii_digit()) => head,
24        _ => p,
25    };
26    Some(if p.starts_with('/') || cwd.is_empty() { p.to_string() } else { format!("{}/{}", cwd.trim_end_matches('/'), p) })
27}
28
29fn classify(name: &str, args: &Value, cwd: &str) -> Vec<Kind> {
30    let s = |k: &str| args.get(k).and_then(Value::as_str);
31    match name {
32        "read" => s("path").and_then(|p| local_path(cwd, p)).map(|path| Kind::Read { path }).into_iter().collect(),
33        "write" => s("path").and_then(|p| local_path(cwd, p)).map(|path| Kind::Edit { path }).into_iter().collect(),
34        "edit" => {
35            let patch = s("input").or_else(|| s("_input")).unwrap_or_default();
36            let mut files: Vec<Kind> = PATCH_FILE
37                .captures_iter(patch)
38                .filter_map(|c| c.get(1).or_else(|| c.get(2)))
39                .filter_map(|m| local_path(cwd, m.as_str().trim()))
40                .map(|path| Kind::Edit { path })
41                .collect();
42            files.dedup();
43            if files.is_empty() { s("path").and_then(|p| local_path(cwd, p)).map(|path| Kind::Edit { path }).into_iter().collect() } else { files }
44        }
45        "bash" => vec![Kind::Bash { cmd: s("command").unwrap_or_default().to_string(), failed: false, error: None, error_line: None }],
46        "grep" | "glob" | "find" | "search" | "lsp" => vec![Kind::Search],
47        n if n.starts_with("mcp__serena_") || n.starts_with("mcp__codebase-memory-mcp_") => vec![Kind::Search],
48        _ => vec![Kind::Other],
49    }
50}
51
52pub fn parse(path: &Path) -> std::io::Result<Session> {
53    let file = std::fs::File::open(path)?;
54    let stem = path.file_stem().map(|s| s.to_string_lossy().into_owned()).unwrap_or_default();
55    let mut sess = Session::new(stem, Harness::Omp);
56    let mut pending: HashMap<String, (i64, Vec<Kind>)> = HashMap::new();
57    let mut narrated = false;
58
59    for line in BufReader::new(file).lines() {
60        let Ok(line) = line else { continue };
61        let Ok(v) = serde_json::from_str::<Value>(&line) else { continue };
62        let ts = v.get("timestamp").and_then(Value::as_str).map_or(0, parse_ts);
63        match v.get("type").and_then(Value::as_str) {
64            Some("session") => {
65                if let Some(c) = v.get("cwd").and_then(Value::as_str) {
66                    sess.cwd = c.to_string();
67                }
68                if let Some(id) = v.get("id").and_then(Value::as_str) {
69                    sess.id = id.to_string();
70                }
71                if v.get("parentSession").is_some() {
72                    sess.sidechain = true;
73                }
74            }
75            Some("message") => {
76                sess.see(ts);
77                let Some(msg) = v.get("message") else { continue };
78                match msg.get("role").and_then(Value::as_str) {
79                    Some("assistant") => {
80                        if let Some(u) = msg.get("usage") {
81                            let n = |k: &str| u.get(k).and_then(Value::as_u64).unwrap_or(0);
82                            sess.input_tokens += n("input") + n("cacheRead") + n("cacheWrite");
83                            sess.output_tokens += n("output");
84                        }
85                        for block in msg.get("content").and_then(Value::as_array).into_iter().flatten() {
86                            match block.get("type").and_then(Value::as_str) {
87                                Some("thinking") => {
88                                    sess.reasoning.thinking_blocks += 1;
89                                    if block.get("thinking").and_then(Value::as_str).is_some_and(|t| !t.trim().is_empty()) {
90                                        sess.reasoning.thinking_nonempty += 1;
91                                    }
92                                }
93                                Some("text") => {
94                                    narrated |= block.get("text").and_then(Value::as_str).is_some_and(|t| t.trim().len() > 10);
95                                }
96                                Some("toolCall") => {
97                                    let name = block.get("name").and_then(Value::as_str).unwrap_or_default();
98                                    let args = block.get("arguments").unwrap_or(&Value::Null);
99                                    let kinds = classify(name, args, &sess.cwd);
100                                    if kinds.iter().any(|k| matches!(k, Kind::Edit { .. })) {
101                                        sess.reasoning.edits += 1;
102                                        // omp tools carry a one-line intent (`i`) on every call: free provenance.
103                                        let intent = args.get("i").and_then(Value::as_str).is_some_and(|t| t.len() > 5);
104                                        sess.reasoning.edits_with_text += (narrated || intent) as u32;
105                                    }
106                                    narrated = false;
107                                    if let Some(id) = block.get("id").and_then(Value::as_str) {
108                                        pending.insert(id.to_string(), (ts, kinds));
109                                    }
110                                }
111                                _ => {}
112                            }
113                        }
114                    }
115                    Some("toolResult") => {
116                        let Some((cts, kinds)) = msg.get("toolCallId").and_then(Value::as_str).and_then(|id| pending.remove(id)) else {
117                            continue;
118                        };
119                        let content = msg.get("content").unwrap_or(&Value::Null);
120                        let failed = msg.get("isError").and_then(Value::as_bool) == Some(true);
121                        let n = kinds.len().max(1) as u64;
122                        let tokens = approx_tokens(content) / n;
123                        for kind in kinds {
124                            match kind {
125                                Kind::Bash { cmd, .. } => sess.push_bash(cts.max(ts), cmd, failed, content),
126                                kind => sess.events.push(Event { ts: cts.max(ts), kind, result_tokens: tokens }),
127                            }
128                        }
129                    }
130                    _ => {}
131                }
132            }
133            _ => {}
134        }
135    }
136    Ok(sess)
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn patch_files() {
145        let args = serde_json::json!({"input": "*** Begin Patch\n[src/a.rs#06FE]\nPUT 1:\n+x\n[local://plan.md#01]\n*** Update File: b/c.rs\n"});
146        let kinds = classify("edit", &args, "/r");
147        assert_eq!(kinds, vec![Kind::Edit { path: "/r/src/a.rs".into() }, Kind::Edit { path: "/r/plan.md".into() }, Kind::Edit { path: "/r/b/c.rs".into() },]);
148    }
149
150    #[test]
151    fn read_line_selector() {
152        assert_eq!(local_path("/r", "src/x.rs:10-20").as_deref(), Some("/r/src/x.rs"));
153        assert_eq!(local_path("/r", "skill://codebase-memory"), None);
154    }
155}