roboticus_agent/

tool_output_filter.rs

//! Tool output noise filter — cleans tool output before it reaches the LLM.
//!
//! Applied *after* tool execution and truncation, *before* the observation is
//! formatted into the ReAct loop. Reduces token waste from ANSI escape codes,
//! progress bars, duplicate lines, and excessive whitespace.

/// A single filter that transforms tool output text.
pub trait ToolOutputFilter: Send + Sync {
    fn name(&self) -> &str;
    fn filter(&self, tool_name: &str, output: &str) -> String;
}

/// Ordered chain of filters applied sequentially.
pub struct ToolOutputFilterChain {
    filters: Vec<Box<dyn ToolOutputFilter>>,
}

impl ToolOutputFilterChain {
    /// Construct the default filter chain used in the ReAct loop.
    pub fn default_chain() -> Self {
        Self {
            filters: vec![
                Box::new(AnsiStripper),
                Box::new(ProgressLineFilter),
                Box::new(DuplicateLineDeduper),
                Box::new(WhitespaceNormalizer),
            ],
        }
    }

    /// Apply all filters in order, returning the cleaned output.
    pub fn apply(&self, tool_name: &str, output: &str) -> String {
        let mut buf = output.to_string();
        for f in &self.filters {
            buf = f.filter(tool_name, &buf);
        }
        buf
    }
}

// ── Filter implementations ──────────────────────────────────────────────────

/// Strip ANSI escape sequences (SGR, cursor, erase, etc.).
pub struct AnsiStripper;

impl ToolOutputFilter for AnsiStripper {
    fn name(&self) -> &str {
        "ansi_stripper"
    }

    fn filter(&self, _tool_name: &str, output: &str) -> String {
        // Match ESC [ ... final_byte  and  ESC ] ... ST  (OSC sequences)
        let mut result = String::with_capacity(output.len());
        let mut chars = output.chars().peekable();
        while let Some(ch) = chars.next() {
            if ch == '\x1b' {
                // CSI sequence: ESC [
                if chars.peek() == Some(&'[') {
                    chars.next(); // consume '['
                    // consume until final byte (0x40-0x7E)
                    while let Some(&c) = chars.peek() {
                        chars.next();
                        if c.is_ascii() && (0x40..=0x7E).contains(&(c as u8)) {
                            break;
                        }
                    }
                    continue;
                }
                // OSC sequence: ESC ]
                if chars.peek() == Some(&']') {
                    chars.next();
                    // consume until ST (ESC \) or BEL (\x07)
                    while let Some(c) = chars.next() {
                        if c == '\x07' {
                            break;
                        }
                        if c == '\x1b' && chars.peek() == Some(&'\\') {
                            chars.next();
                            break;
                        }
                    }
                    continue;
                }
                // Single-char escape (ESC + one byte)
                chars.next();
                continue;
            }
            // Strip carriage return (common in progress output)
            if ch == '\r' {
                continue;
            }
            result.push(ch);
        }
        result
    }
}

/// Remove lines that look like progress bars or spinners.
pub struct ProgressLineFilter;

impl ToolOutputFilter for ProgressLineFilter {
    fn name(&self) -> &str {
        "progress_line_filter"
    }

    fn filter(&self, _tool_name: &str, output: &str) -> String {
        output
            .lines()
            .filter(|line| {
                let trimmed = line.trim();
                // Skip lines that are mostly progress bar characters
                if trimmed.is_empty() {
                    return true;
                }
                let progress_chars = trimmed
                    .chars()
                    .filter(|c| {
                        matches!(
                            c,
                            '█' | '▓'
                                | '▒'
                                | '░'
                                | '━'
                                | '─'
                                | '┃'
                                | '│'
                                | '⠋'
                                | '⠙'
                                | '⠹'
                                | '⠸'
                                | '⠼'
                                | '⠴'
                                | '⠦'
                                | '⠧'
                                | '⠇'
                                | '⠏'
                                | '/'
                                | '\\'
                                | '-'
                                | '='
                                | '>'
                                | '#'
                                | '.'
                                | '●'
                                | '○'
                        )
                    })
                    .count();
                let total = trimmed.chars().count();
                // If >60% of the line is progress characters and it contains '%', skip it
                if total > 5 && progress_chars * 100 / total > 60 {
                    return false;
                }
                // Skip lines that are just a percentage update
                if trimmed
                    .chars()
                    .all(|c| c.is_ascii_digit() || c == '%' || c == '.' || c == ' ')
                    && trimmed.contains('%')
                {
                    return false;
                }
                true
            })
            .collect::<Vec<_>>()
            .join("\n")
    }
}

/// Remove consecutive duplicate lines (common in build/test output).
pub struct DuplicateLineDeduper;

impl ToolOutputFilter for DuplicateLineDeduper {
    fn name(&self) -> &str {
        "duplicate_line_deduper"
    }

    fn filter(&self, _tool_name: &str, output: &str) -> String {
        let mut result = Vec::new();
        let mut prev: Option<&str> = None;
        let mut dup_count = 0u32;
        for line in output.lines() {
            if Some(line) == prev {
                dup_count += 1;
                continue;
            }
            if dup_count > 0 {
                result.push(format!("  [... repeated {dup_count} more time(s)]"));
                dup_count = 0;
            }
            result.push(line.to_string());
            prev = Some(line);
        }
        if dup_count > 0 {
            result.push(format!("  [... repeated {dup_count} more time(s)]"));
        }
        result.join("\n")
    }
}

/// Collapse runs of blank lines to max 2, trim trailing whitespace per line.
pub struct WhitespaceNormalizer;

impl ToolOutputFilter for WhitespaceNormalizer {
    fn name(&self) -> &str {
        "whitespace_normalizer"
    }

    fn filter(&self, _tool_name: &str, output: &str) -> String {
        let mut result = Vec::new();
        let mut blank_run = 0u32;
        for line in output.lines() {
            let trimmed = line.trim_end();
            if trimmed.is_empty() {
                blank_run += 1;
                if blank_run <= 2 {
                    result.push(String::new());
                }
            } else {
                blank_run = 0;
                result.push(trimmed.to_string());
            }
        }
        // Trim trailing blank lines
        while result.last().is_some_and(|l| l.is_empty()) {
            result.pop();
        }
        result.join("\n")
    }
}

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

    #[test]
    fn ansi_stripper_removes_sgr() {
        let input = "\x1b[31mERROR\x1b[0m: something failed";
        let chain = ToolOutputFilterChain::default_chain();
        let out = chain.apply("test", input);
        assert_eq!(out, "ERROR: something failed");
    }

    #[test]
    fn progress_filter_removes_bars() {
        let input = "Compiling foo\n███████████████░░░░░ 75%\nDone.";
        let filter = ProgressLineFilter;
        let out = filter.filter("test", input);
        assert!(out.contains("Compiling foo"));
        assert!(out.contains("Done."));
        assert!(!out.contains("███"));
    }

    #[test]
    fn deduper_collapses_repeats() {
        let input = "line1\nline2\nline2\nline2\nline3";
        let filter = DuplicateLineDeduper;
        let out = filter.filter("test", input);
        assert!(out.contains("line2"));
        assert!(out.contains("[... repeated 2 more time(s)]"));
        assert!(out.contains("line3"));
    }

    #[test]
    fn whitespace_normalizer_collapses_blanks() {
        let input = "a\n\n\n\n\nb\n\nc";
        let filter = WhitespaceNormalizer;
        let out = filter.filter("test", input);
        // Max 2 blank lines between a and b
        assert_eq!(out, "a\n\n\nb\n\nc");
    }

    #[test]
    fn full_chain_applies_all() {
        let input = "\x1b[32mok\x1b[0m\nfoo\nfoo\nfoo\n\n\n\n\nbar";
        let chain = ToolOutputFilterChain::default_chain();
        let out = chain.apply("test", input);
        assert!(out.starts_with("ok"));
        assert!(out.contains("[... repeated 2 more time(s)]"));
        assert!(out.contains("bar"));
    }
}