utils/
tool_result_truncator.rs

1use crate::temp_dir::TempDir;
2use std::path::Path;
3
4#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5pub struct ToolResultTruncator {
6    pub head: usize,
7    pub tail: usize,
8}
9
10impl ToolResultTruncator {
11    pub fn truncate(&self, text: String, dir: &TempDir, label: &str) -> String {
12        let Some((head, tail)) = self.split(&text) else {
13            return text;
14        };
15        let saved = match dir.save(label, &text) {
16            Ok(path) => format!("{SAVED_TO}{}", path.display()),
17            Err(error) => {
18                tracing::warn!("Failed to save oversized output: {error}");
19                "the full output could not be saved".to_string()
20            }
21        };
22        let omitted = text.len() - head.len() - tail.len();
23        format!("{head}\n[... {omitted} bytes omitted; {saved}{MARKER_END}{tail}")
24    }
25
26    fn split<'a>(&self, text: &'a str) -> Option<(&'a str, &'a str)> {
27        (text.len() > self.head + self.tail).then(|| {
28            let head = &text[..text.floor_char_boundary(self.head)];
29            let tail = &text[text.ceil_char_boundary(text.len() - self.tail)..];
30            (head, tail)
31        })
32    }
33}
34
35pub fn saved_path(excerpt: &str) -> Option<&Path> {
36    let (_, rest) = excerpt.split_once(SAVED_TO)?;
37    rest.split_once(MARKER_END).map(|(path, _)| Path::new(path))
38}
39
40const SAVED_TO: &str = "full output saved to ";
41const MARKER_END: &str = " ...]\n";
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46    use std::fs::read_to_string;
47
48    const TRUNCATOR: ToolResultTruncator = ToolResultTruncator { head: 16, tail: 12 };
49
50    fn truncated(text: &str) -> (String, TempDir) {
51        let dir = TempDir::new();
52        (TRUNCATOR.truncate(text.to_string(), &dir, "out"), dir)
53    }
54
55    fn head_and_tail(excerpt: &str) -> (&str, &str) {
56        let (head, rest) = excerpt.split_once("\n[... ").unwrap();
57        (head, rest.split_once(MARKER_END).unwrap().1)
58    }
59
60    #[test]
61    fn text_is_truncated_once_it_outgrows_the_head_and_tail() {
62        let (fits, _dir) = truncated(&"x".repeat(28));
63        assert_eq!(fits, "x".repeat(28));
64        assert_eq!(saved_path(&fits), None);
65
66        let (excerpt, _dir) = truncated(&"x".repeat(29));
67        assert_eq!(head_and_tail(&excerpt), ("x".repeat(16).as_str(), "x".repeat(12).as_str()));
68    }
69
70    #[test]
71    fn truncated_text_keeps_byte_limited_excerpts_and_saves_the_full_text() {
72        let text = (1..=20).map(|n| format!("line {n}\n")).collect::<Vec<_>>().concat();
73
74        let (excerpt, _dir) = truncated(&text);
75
76        let saved = saved_path(&excerpt).unwrap();
77        let head = "line 1\nline 2\nli";
78        let tail = " 19\nline 20\n";
79        let omitted = text.len() - head.len() - tail.len();
80        let marker = format!("[... {omitted} bytes omitted; full output saved to {} ...]", saved.display());
81        assert_eq!(excerpt, format!("{head}\n{marker}\n{tail}"));
82        assert_eq!(read_to_string(saved).unwrap(), text);
83    }
84
85    #[test]
86    fn truncate_keeps_the_final_line_even_when_it_overflows_the_tail() {
87        let (excerpt, _dir) = truncated(&format!("short\n{}\n", "x".repeat(50)));
88        assert_eq!(head_and_tail(&excerpt).1, format!("{}\n", "x".repeat(11)));
89    }
90
91    #[test]
92    fn excerpts_stay_within_their_budgets_at_utf8_boundaries() {
93        let (excerpt, _dir) = truncated(&"🦀".repeat(20));
94        assert_eq!(head_and_tail(&excerpt), ("🦀".repeat(4).as_str(), "🦀".repeat(3).as_str()));
95
96        let (excerpt, _dir) = truncated(&"─".repeat(20));
97        assert_eq!(head_and_tail(&excerpt), ("─".repeat(5).as_str(), "─".repeat(4).as_str()));
98    }
99}
utils/tool_result_truncator.rs

utils/
tool_result_truncator.rs
