aether-utils 0.2.18

Shared utilities for the Aether AI agent framework
Documentation
use crate::temp_dir::TempDir;
use std::path::Path;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ToolResultTruncator {
    pub head: usize,
    pub tail: usize,
}

impl ToolResultTruncator {
    pub fn truncate(&self, text: String, dir: &TempDir, label: &str) -> String {
        let Some((head, tail)) = self.split(&text) else {
            return text;
        };
        let saved = match dir.save(label, &text) {
            Ok(path) => format!("{SAVED_TO}{}", path.display()),
            Err(error) => {
                tracing::warn!("Failed to save oversized output: {error}");
                "the full output could not be saved".to_string()
            }
        };
        let omitted = text.len() - head.len() - tail.len();
        format!("{head}\n[... {omitted} bytes omitted; {saved}{MARKER_END}{tail}")
    }

    fn split<'a>(&self, text: &'a str) -> Option<(&'a str, &'a str)> {
        (text.len() > self.head + self.tail).then(|| {
            let head = &text[..text.floor_char_boundary(self.head)];
            let tail = &text[text.ceil_char_boundary(text.len() - self.tail)..];
            (head, tail)
        })
    }
}

pub fn saved_path(excerpt: &str) -> Option<&Path> {
    let (_, rest) = excerpt.split_once(SAVED_TO)?;
    rest.split_once(MARKER_END).map(|(path, _)| Path::new(path))
}

const SAVED_TO: &str = "full output saved to ";
const MARKER_END: &str = " ...]\n";

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs::read_to_string;

    const TRUNCATOR: ToolResultTruncator = ToolResultTruncator { head: 16, tail: 12 };

    fn truncated(text: &str) -> (String, TempDir) {
        let dir = TempDir::new();
        (TRUNCATOR.truncate(text.to_string(), &dir, "out"), dir)
    }

    fn head_and_tail(excerpt: &str) -> (&str, &str) {
        let (head, rest) = excerpt.split_once("\n[... ").unwrap();
        (head, rest.split_once(MARKER_END).unwrap().1)
    }

    #[test]
    fn text_is_truncated_once_it_outgrows_the_head_and_tail() {
        let (fits, _dir) = truncated(&"x".repeat(28));
        assert_eq!(fits, "x".repeat(28));
        assert_eq!(saved_path(&fits), None);

        let (excerpt, _dir) = truncated(&"x".repeat(29));
        assert_eq!(head_and_tail(&excerpt), ("x".repeat(16).as_str(), "x".repeat(12).as_str()));
    }

    #[test]
    fn truncated_text_keeps_byte_limited_excerpts_and_saves_the_full_text() {
        let text = (1..=20).map(|n| format!("line {n}\n")).collect::<Vec<_>>().concat();

        let (excerpt, _dir) = truncated(&text);

        let saved = saved_path(&excerpt).unwrap();
        let head = "line 1\nline 2\nli";
        let tail = " 19\nline 20\n";
        let omitted = text.len() - head.len() - tail.len();
        let marker = format!("[... {omitted} bytes omitted; full output saved to {} ...]", saved.display());
        assert_eq!(excerpt, format!("{head}\n{marker}\n{tail}"));
        assert_eq!(read_to_string(saved).unwrap(), text);
    }

    #[test]
    fn truncate_keeps_the_final_line_even_when_it_overflows_the_tail() {
        let (excerpt, _dir) = truncated(&format!("short\n{}\n", "x".repeat(50)));
        assert_eq!(head_and_tail(&excerpt).1, format!("{}\n", "x".repeat(11)));
    }

    #[test]
    fn excerpts_stay_within_their_budgets_at_utf8_boundaries() {
        let (excerpt, _dir) = truncated(&"🦀".repeat(20));
        assert_eq!(head_and_tail(&excerpt), ("🦀".repeat(4).as_str(), "🦀".repeat(3).as_str()));

        let (excerpt, _dir) = truncated(&"─".repeat(20));
        assert_eq!(head_and_tail(&excerpt), ("─".repeat(5).as_str(), "─".repeat(4).as_str()));
    }
}