agx_core/
codex.rs

1use crate::timeline::{
2    self, Step, Usage, assistant_text_step, attach_usage_to_first, compute_durations, parse_iso_ms,
3    pretty_json, tool_result_step, tool_use_step, user_text_step,
4};
5use anyhow::{Context, Result};
6use serde::Deserialize;
7use std::collections::HashMap;
8use std::path::Path;
9
10#[derive(Debug, Deserialize)]
11struct Entry {
12    #[serde(default)]
13    timestamp: Option<String>,
14    #[serde(rename = "type")]
15    kind: String,
16    #[serde(default)]
17    payload: serde_json::Value,
18}
19
20pub fn load(path: &Path) -> Result<Vec<Step>> {
21    // Line-stream via BufReader (same rationale as session.rs). Codex
22    // rollouts in the wild routinely grow past 10MB of JSONL; holding
23    // the whole file as a `String` just to call `.lines()` on it was
24    // pure waste.
25    use std::fs::File;
26    use std::io::{BufRead, BufReader};
27    let file = File::open(path)
28        .with_context(|| format!("opening codex session file: {}", path.display()))?;
29    let reader = BufReader::new(file);
30    let mut entries: Vec<Entry> = Vec::with_capacity(1024);
31    for (i, line) in reader.lines().enumerate() {
32        let line = line.with_context(|| format!("reading line {} of codex session", i + 1))?;
33        if line.trim().is_empty() {
34            continue;
35        }
36        let entry = serde_json::from_str::<Entry>(&line)
37            .with_context(|| format!("parsing line {} of codex session", i + 1))?;
38        entries.push(entry);
39    }
40
41    let tool_meta = collect_tool_meta(&entries);
42    let mut steps = Vec::new();
43    for entry in &entries {
44        if entry.kind != "response_item" {
45            continue;
46        }
47        let payload_type = entry.payload.get("type").and_then(|t| t.as_str());
48        let ts = entry.timestamp.as_deref().and_then(parse_iso_ms);
49        let mut maybe_step: Option<Step> = None;
50        match payload_type {
51            Some("message") => {
52                let role = entry
53                    .payload
54                    .get("role")
55                    .and_then(|v| v.as_str())
56                    .unwrap_or("");
57                let text = extract_message_text(&entry.payload);
58                if !text.trim().is_empty() {
59                    maybe_step = match role {
60                        "user" => Some(user_text_step(&text)),
61                        "assistant" => Some(assistant_text_step(&text)),
62                        _ => None,
63                    };
64                }
65            }
66            Some("function_call") => {
67                let call_id = entry
68                    .payload
69                    .get("call_id")
70                    .and_then(|v| v.as_str())
71                    .unwrap_or("");
72                let name = entry
73                    .payload
74                    .get("name")
75                    .and_then(|v| v.as_str())
76                    .unwrap_or("(unknown)");
77                let input_pretty = prettify_codex_arguments(&entry.payload);
78                maybe_step = Some(tool_use_step(call_id, name, &input_pretty));
79            }
80            Some("function_call_output") => {
81                let call_id = entry
82                    .payload
83                    .get("call_id")
84                    .and_then(|v| v.as_str())
85                    .unwrap_or("");
86                let output = entry
87                    .payload
88                    .get("output")
89                    .and_then(|v| v.as_str())
90                    .map_or_else(|| pretty_json(&entry.payload.get("output")), String::from);
91                let meta = tool_meta.get(call_id);
92                maybe_step = Some(tool_result_step(
93                    call_id,
94                    &output,
95                    meta.map(|m| m.name.as_str()),
96                    meta.map(|m| m.input_pretty.as_str()),
97                ));
98            }
99            _ => {}
100        }
101        if let Some(mut step) = maybe_step {
102            step.timestamp_ms = ts;
103            let is_assistant_message = payload_type == Some("message")
104                && entry.payload.get("role").and_then(|v| v.as_str()) == Some("assistant");
105            steps.push(step);
106            if is_assistant_message {
107                let idx = steps.len() - 1;
108                let model = entry.payload.get("model").and_then(|v| v.as_str());
109                let usage = extract_codex_usage(&entry.payload);
110                attach_usage_to_first(&mut steps, idx, model, &usage);
111            }
112        }
113    }
114    compute_durations(&mut steps);
115    Ok(steps)
116}
117
118/// Codex payload usage shape mirrors OpenAI Responses API conventions.
119/// Accepts either snake_case (`input_tokens`) or legacy camelCase
120/// (`promptTokens`) — Codex has used both across its versions.
121fn extract_codex_usage(payload: &serde_json::Value) -> Usage {
122    let Some(usage_obj) = payload.get("usage") else {
123        return Usage::default();
124    };
125    let get_u64 = |keys: &[&str]| -> Option<u64> {
126        for k in keys {
127            if let Some(v) = usage_obj.get(*k).and_then(|v| v.as_u64()) {
128                return Some(v);
129            }
130        }
131        None
132    };
133    Usage {
134        tokens_in: get_u64(&["input_tokens", "prompt_tokens", "promptTokens"]),
135        tokens_out: get_u64(&["output_tokens", "completion_tokens", "completionTokens"]),
136        cache_read: get_u64(&["cached_tokens", "prompt_cache_read_tokens"]),
137        cache_create: None,
138    }
139}
140
141#[derive(Debug, Clone)]
142struct ToolMeta {
143    name: String,
144    input_pretty: String,
145}
146
147fn collect_tool_meta(entries: &[Entry]) -> HashMap<String, ToolMeta> {
148    let mut map = HashMap::new();
149    for entry in entries {
150        if entry.kind != "response_item" {
151            continue;
152        }
153        if entry.payload.get("type").and_then(|t| t.as_str()) != Some("function_call") {
154            continue;
155        }
156        let Some(call_id) = entry.payload.get("call_id").and_then(|v| v.as_str()) else {
157            continue;
158        };
159        let name = entry
160            .payload
161            .get("name")
162            .and_then(|v| v.as_str())
163            .unwrap_or("(unknown)")
164            .to_string();
165        let input_pretty = prettify_codex_arguments(&entry.payload);
166        map.insert(call_id.to_string(), ToolMeta { name, input_pretty });
167    }
168    map
169}
170
171// Codex stores function_call arguments as a serialized JSON string inside
172// the `arguments` field. Try to re-parse and pretty-print; fall back to the
173// raw string if that fails.
174fn prettify_codex_arguments(payload: &serde_json::Value) -> String {
175    let raw = payload
176        .get("arguments")
177        .and_then(|v| v.as_str())
178        .unwrap_or("");
179    if raw.is_empty() {
180        return String::new();
181    }
182    match serde_json::from_str::<serde_json::Value>(raw) {
183        Ok(v) => timeline::pretty_json(&v),
184        Err(_) => raw.to_string(),
185    }
186}
187
188// A codex message payload has `content: [{type: input_text|output_text, text: "..."}, ...]`.
189// Concatenate all text fragments into one string for the step detail.
190fn extract_message_text(payload: &serde_json::Value) -> String {
191    let Some(items) = payload.get("content").and_then(|c| c.as_array()) else {
192        return String::new();
193    };
194    items
195        .iter()
196        .filter_map(|item| item.get("text").and_then(|t| t.as_str()))
197        .collect::<Vec<_>>()
198        .join("\n")
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204    use crate::timeline::StepKind;
205    use std::io::Write;
206    use tempfile::NamedTempFile;
207
208    fn write_file(content: &str) -> NamedTempFile {
209        let mut f = NamedTempFile::new().unwrap();
210        f.write_all(content.as_bytes()).unwrap();
211        f
212    }
213
214    #[test]
215    fn parses_user_and_assistant_messages() {
216        let jsonl = r#"{"timestamp":"2024-01-01T00:00:00Z","type":"session_meta","payload":{"id":"s1","cwd":"/tmp"}}
217{"timestamp":"2024-01-01T00:00:01Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}}
218{"timestamp":"2024-01-01T00:00:02Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"hi there"}]}}
219"#;
220        let f = write_file(jsonl);
221        let steps = load(f.path()).unwrap();
222        assert_eq!(steps.len(), 2);
223        assert_eq!(steps[0].kind, StepKind::UserText);
224        assert!(steps[0].detail.contains("hello"));
225        assert_eq!(steps[1].kind, StepKind::AssistantText);
226        assert!(steps[1].detail.contains("hi there"));
227    }
228
229    #[test]
230    fn pairs_function_call_with_function_call_output() {
231        let jsonl = r#"{"timestamp":"2024-01-01T00:00:00Z","type":"response_item","payload":{"type":"function_call","call_id":"call_abc","name":"exec_command","arguments":"{\"cmd\":\"ls\"}"}}
232{"timestamp":"2024-01-01T00:00:01Z","type":"response_item","payload":{"type":"function_call_output","call_id":"call_abc","output":"file1\nfile2"}}
233"#;
234        let f = write_file(jsonl);
235        let steps = load(f.path()).unwrap();
236        assert_eq!(steps.len(), 2);
237        assert_eq!(steps[0].kind, StepKind::ToolUse);
238        assert!(steps[0].detail.contains("exec_command"));
239        assert!(steps[0].detail.contains("\"cmd\""));
240        assert!(steps[0].detail.contains("\"ls\""));
241        assert_eq!(steps[1].kind, StepKind::ToolResult);
242        assert!(steps[1].label.contains("exec_command"));
243        assert!(steps[1].detail.contains("Tool: exec_command"));
244        assert!(steps[1].detail.contains("Input:"));
245        assert!(steps[1].detail.contains("Result:"));
246        assert!(steps[1].detail.contains("file1"));
247    }
248
249    #[test]
250    fn skips_developer_role_messages() {
251        let jsonl = r#"{"timestamp":"2024-01-01T00:00:00Z","type":"response_item","payload":{"type":"message","role":"developer","content":[{"type":"input_text","text":"system policies..."}]}}
252{"timestamp":"2024-01-01T00:00:01Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"real question"}]}}
253"#;
254        let f = write_file(jsonl);
255        let steps = load(f.path()).unwrap();
256        assert_eq!(steps.len(), 1);
257        assert_eq!(steps[0].kind, StepKind::UserText);
258        assert!(steps[0].detail.contains("real question"));
259    }
260
261    #[test]
262    fn skips_reasoning_entries() {
263        let jsonl = r#"{"timestamp":"2024-01-01T00:00:00Z","type":"response_item","payload":{"type":"reasoning","summary":[],"content":null}}
264{"timestamp":"2024-01-01T00:00:01Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}]}}
265"#;
266        let f = write_file(jsonl);
267        let steps = load(f.path()).unwrap();
268        assert_eq!(steps.len(), 1);
269        assert_eq!(steps[0].kind, StepKind::AssistantText);
270    }
271
272    #[test]
273    fn parses_usage_and_model_on_assistant_message() {
274        let jsonl = r#"{"timestamp":"2024-01-01T00:00:00Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"answer"}],"model":"gpt-5","usage":{"input_tokens":120,"output_tokens":60,"cached_tokens":40}}}
275"#;
276        let f = write_file(jsonl);
277        let steps = load(f.path()).unwrap();
278        assert_eq!(steps.len(), 1);
279        assert_eq!(steps[0].model.as_deref(), Some("gpt-5"));
280        assert_eq!(steps[0].tokens_in, Some(120));
281        assert_eq!(steps[0].tokens_out, Some(60));
282        assert_eq!(steps[0].cache_read, Some(40));
283    }
284
285    #[test]
286    fn legacy_camelcase_usage_fields_parse() {
287        // Older Codex rollouts used `promptTokens` / `completionTokens` —
288        // cover the fallback path in extract_codex_usage.
289        let jsonl = r#"{"timestamp":"2024-01-01T00:00:00Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"x"}],"usage":{"promptTokens":10,"completionTokens":20}}}
290"#;
291        let f = write_file(jsonl);
292        let steps = load(f.path()).unwrap();
293        assert_eq!(steps[0].tokens_in, Some(10));
294        assert_eq!(steps[0].tokens_out, Some(20));
295    }
296
297    #[test]
298    fn user_message_does_not_get_usage() {
299        // Usage attaches only to assistant messages; a user message with a
300        // usage field (unusual but not impossible) should be ignored.
301        let jsonl = r#"{"timestamp":"2024-01-01T00:00:00Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}],"usage":{"input_tokens":5}}}
302"#;
303        let f = write_file(jsonl);
304        let steps = load(f.path()).unwrap();
305        assert_eq!(steps[0].tokens_in, None);
306    }
307
308    #[test]
309    fn skips_non_response_item_entries() {
310        let jsonl = r#"{"timestamp":"2024-01-01T00:00:00Z","type":"session_meta","payload":{"id":"s1"}}
311{"timestamp":"2024-01-01T00:00:01Z","type":"event_msg","payload":{"type":"task_started"}}
312{"timestamp":"2024-01-01T00:00:02Z","type":"turn_context","payload":{}}
313{"timestamp":"2024-01-01T00:00:03Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]}}
314"#;
315        let f = write_file(jsonl);
316        let steps = load(f.path()).unwrap();
317        assert_eq!(steps.len(), 1);
318        assert_eq!(steps[0].kind, StepKind::UserText);
319    }
320}
agx_core/codex.rs

agx_core/
codex.rs
