Skip to main content

edda_transcript/
extract.rs

1use std::io::{BufRead, BufReader};
2use std::path::Path;
3
4/// Extract the last assistant text message from a stored transcript JSONL.
5///
6/// Scans the file line by line (forward), keeping track of the last assistant
7/// text seen. Returns the final one, truncated to `max_chars`.
8///
9/// Expected format per line:
10/// ```json
11/// {"type":"assistant","message":{"content":[{"type":"text","text":"..."},{"type":"tool_use",...}]}}
12/// ```
13///
14/// Only extracts `content` blocks with `"type":"text"`. Skips `tool_use` blocks.
15pub fn extract_last_assistant_text(store_path: &Path, max_chars: usize) -> Option<String> {
16    let file = std::fs::File::open(store_path).ok()?;
17    let reader = BufReader::new(file);
18
19    let mut last_text: Option<String> = None;
20
21    for line in reader.lines() {
22        let line = line.ok()?;
23        if line.is_empty() {
24            continue;
25        }
26
27        let parsed: serde_json::Value = match serde_json::from_str(&line) {
28            Ok(v) => v,
29            Err(_) => continue,
30        };
31
32        if parsed.get("type").and_then(|v| v.as_str()) != Some("assistant") {
33            continue;
34        }
35
36        // Extract text from message.content array
37        if let Some(content_arr) = parsed
38            .get("message")
39            .and_then(|m| m.get("content"))
40            .and_then(|c| c.as_array())
41        {
42            let mut texts: Vec<&str> = Vec::new();
43            for block in content_arr {
44                if block.get("type").and_then(|t| t.as_str()) == Some("text") {
45                    if let Some(text) = block.get("text").and_then(|t| t.as_str()) {
46                        texts.push(text);
47                    }
48                }
49            }
50            if !texts.is_empty() {
51                last_text = Some(texts.join("\n"));
52            }
53        }
54    }
55
56    // Truncate to max_chars
57    last_text.map(|t| {
58        if t.len() > max_chars {
59            let truncated = &t[..t.floor_char_boundary(max_chars)];
60            format!("{truncated}...")
61        } else {
62            t
63        }
64    })
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use std::io::Write;
71
72    fn write_store(dir: &Path, lines: &[&str]) -> std::path::PathBuf {
73        let path = dir.join("test-session.jsonl");
74        let mut f = std::fs::File::create(&path).unwrap();
75        for line in lines {
76            writeln!(f, "{line}").unwrap();
77        }
78        path
79    }
80
81    #[test]
82    fn extract_basic_assistant_text() {
83        let tmp = tempfile::tempdir().unwrap();
84        let store = write_store(
85            tmp.path(),
86            &[
87                r#"{"type":"user","uuid":"u1","message":{"content":"hello"}}"#,
88                r#"{"type":"assistant","uuid":"a1","message":{"content":[{"type":"text","text":"Hi there!"}]}}"#,
89            ],
90        );
91
92        let result = extract_last_assistant_text(&store, 500);
93        assert_eq!(result.as_deref(), Some("Hi there!"));
94    }
95
96    #[test]
97    fn extract_skips_tool_use_blocks() {
98        let tmp = tempfile::tempdir().unwrap();
99        let store = write_store(
100            tmp.path(),
101            &[
102                r#"{"type":"assistant","uuid":"a1","message":{"content":[{"type":"text","text":"Let me check"},{"type":"tool_use","id":"tu1","name":"Bash"}]}}"#,
103            ],
104        );
105
106        let result = extract_last_assistant_text(&store, 500);
107        assert_eq!(result.as_deref(), Some("Let me check"));
108    }
109
110    #[test]
111    fn extract_returns_last_assistant() {
112        let tmp = tempfile::tempdir().unwrap();
113        let store = write_store(
114            tmp.path(),
115            &[
116                r#"{"type":"assistant","uuid":"a1","message":{"content":[{"type":"text","text":"first"}]}}"#,
117                r#"{"type":"user","uuid":"u2","message":{"content":"more"}}"#,
118                r#"{"type":"assistant","uuid":"a2","message":{"content":[{"type":"text","text":"second and final"}]}}"#,
119            ],
120        );
121
122        let result = extract_last_assistant_text(&store, 500);
123        assert_eq!(result.as_deref(), Some("second and final"));
124    }
125
126    #[test]
127    fn extract_truncates_long_text() {
128        let tmp = tempfile::tempdir().unwrap();
129        let long_text = "x".repeat(1000);
130        let line = format!(
131            r#"{{"type":"assistant","uuid":"a1","message":{{"content":[{{"type":"text","text":"{long_text}"}}]}}}}"#
132        );
133        let store = write_store(tmp.path(), &[&line]);
134
135        let result = extract_last_assistant_text(&store, 100).unwrap();
136        assert!(result.len() <= 104); // 100 chars + "..."
137        assert!(result.ends_with("..."));
138    }
139
140    #[test]
141    fn extract_no_assistant_returns_none() {
142        let tmp = tempfile::tempdir().unwrap();
143        let store = write_store(
144            tmp.path(),
145            &[r#"{"type":"user","uuid":"u1","message":{"content":"hello"}}"#],
146        );
147
148        let result = extract_last_assistant_text(&store, 500);
149        assert!(result.is_none());
150    }
151
152    #[test]
153    fn extract_missing_file_returns_none() {
154        let result = extract_last_assistant_text(Path::new("/nonexistent/file.jsonl"), 500);
155        assert!(result.is_none());
156    }
157}