zesdex/src/app/runtime/context/squash.rs

#![allow(dead_code)]
//! Per-tool-result compression: shrink large tool outputs before they
//! ever enter conversation history, dispatching by content shape.
//!
//! Flow: `apply(tool_name, output)` -> `read` tool or under the size
//! floor? pass through unchanged : valid JSON? `squash_json` : tool is
//! `bash` and looks log-shaped? `squash_log` : `squash_generic`.
//!
//! Why: a single large `bash`/`grep` result can dominate a
//! conversation's token budget even on its first occurrence, long
//! before `dedup`/`shaping` ever get a chance to act on repeats or
//! overall budget.
use std::collections::HashSet;
use std::fmt::Write;

/// Below this size, compression isn't worth the risk of losing detail —
/// pass the output through unchanged.
const SQUASH_FLOOR_BYTES: usize = 1500;

/// Byte budget for the generic fallback compressor — double the squash
/// floor, so the fallback path still yields a real reduction on
/// anything that triggered it.
const GENERIC_BUDGET_BYTES: usize = SQUASH_FLOOR_BYTES * 2;

/// Tools whose output must never be altered. `read` is exempted because
/// its output must stay byte-exact — the agent relies on it for
/// exact-match edits afterward, and squashing a file that happens to
/// parse as JSON (e.g. `package.json`) would silently corrupt the
/// agent's view of real file content.
const NEVER_SQUASH: &[&str] = &["read"];

/// Tools whose output the log classifier is allowed to run on.
/// `looks_log_shaped` keys purely on content (>=3 error/warn/fail-shaped
/// lines), which a `grep`/`search` result full of matches against
/// error-handling code would trip just as easily as a real build log —
/// but `squash_log` caps at 20 error + 10 warning lines with no byte
/// budget, silently dropping legitimate matches past that cap. Only
/// `bash` (the actual log-producing tool) is allowed to route through
/// it; everything else that looks log-shaped falls through to the
/// gentler, byte-budgeted `squash_generic` instead.
const LOG_SHAPED_TOOLS: &[&str] = &["bash"];

/// Compress a tool's raw output before it's stored in conversation
/// history.
///
/// Return: `output` unchanged if `tool_name` is in `NEVER_SQUASH` or at
/// or under `SQUASH_FLOOR_BYTES`; otherwise the compressed form from
/// whichever detector matches its content shape.
pub fn apply(tool_name: &str, output: &str) -> String {
    if NEVER_SQUASH.contains(&tool_name) || output.len() <= SQUASH_FLOOR_BYTES {
        return output.to_string();
    }
    if serde_json::from_str::<serde_json::Value>(output).is_ok() {
        return squash_json(output);
    }
    if LOG_SHAPED_TOOLS.contains(&tool_name) && looks_log_shaped(output) {
        return squash_log(output);
    }
    squash_generic(output, GENERIC_BUDGET_BYTES)
}

/// Compress a JSON tool result by keeping all structural content (keys,
/// array/object shape) and eliding long, low-entropy string *values*,
/// while keeping short values (<=20 chars) and high-entropy single-token
/// ones (UUIDs, hashes, paths) intact. Array elements past the first 3
/// are elided regardless of length/entropy.
///
/// Why walk a parsed `Value` instead of hand-rolling a JSON tokenizer:
/// `serde_json` already handles escaping/nesting correctly (this
/// codebase's own `dto::chat::tool::repair_json` exists specifically to
/// work around how easy it is to get that wrong by hand) — reusing it
/// is both simpler and more robust.
///
/// Return: re-serialized JSON with the same shape as the input.
fn squash_json(text: &str) -> String {
    let Ok(mut value) = serde_json::from_str::<serde_json::Value>(text) else {
        return text.to_string();
    };
    squash_json_value(&mut value, false);
    serde_json::to_string(&value).unwrap_or_else(|_| text.to_string())
}

/// Recursively elide long, low-entropy string values in place.
/// `in_late_array` is true once past the first 3 elements of an
/// enclosing array, tightening the elision rule for the rest of it.
///
/// Why the `!s.contains(' ')` gate before the entropy check: raw
/// per-character Shannon entropy alone does NOT separate "meaningful
/// prose" from "random-looking identifier" — verified empirically,
/// repeated English prose scores ~3.89 bits/char, *higher* than a UUID's
/// ~3.39 or a SHA-256 hex digest's ~3.66, because prose draws from a
/// wide, fairly-balanced character set too. What actually distinguishes
/// identifiers from prose is that identifiers are a single unbroken
/// token — this mirrors headroom's own approach (its entropy check is
/// "cheaply pre-filtered by 'no spaces'" before scoring). Multi-word
/// values never reach the entropy branch at all; only whitespace-free
/// tokens do, where entropy correctly separates "abc123" or "aaaaaaaa"
/// (low, elided if long) from a UUID/hash/API-key-shaped string (high,
/// kept).
fn squash_json_value(value: &mut serde_json::Value, in_late_array: bool) {
    match value {
        serde_json::Value::String(s) => {
            let looks_like_identifier = !s.contains(' ') && shannon_entropy(s) >= 3.0;
            let keep = !in_late_array && (s.len() <= 20 || looks_like_identifier);
            if !keep {
                *s = "…".to_string();
            }
        }
        serde_json::Value::Array(items) => {
            for (i, item) in items.iter_mut().enumerate() {
                squash_json_value(item, i >= 3);
            }
        }
        serde_json::Value::Object(map) => {
            for v in map.values_mut() {
                squash_json_value(v, false);
            }
        }
        _ => {}
    }
}

/// Shannon entropy in bits per character — used, after the `squash_json`
/// caller's own "no internal whitespace" pre-filter, to distinguish
/// high-entropy single-token strings (UUIDs, hashes, random IDs, worth
/// keeping) from low-entropy ones (e.g. `"aaaaaaaaaa"`, safe to elide).
/// 3.0 sits comfortably below a UUID's ~3.39 and a SHA-256 hex digest's
/// ~3.66 (both empirically measured with this exact formula) while
/// staying well above a degenerate repeated-character string's 0.0.
fn shannon_entropy(s: &str) -> f64 {
    if s.is_empty() {
        return 0.0;
    }
    let mut counts: std::collections::HashMap<char, usize> = std::collections::HashMap::new();
    for c in s.chars() {
        *counts.entry(c).or_insert(0) += 1;
    }
    let len = s.chars().count() as f64;
    counts
        .values()
        .map(|&count| {
            let p = f64::from(u32::try_from(count).unwrap_or(u32::MAX)) / len;
            -p * p.log2()
        })
        .sum()
}

/// Coarse severity classification for a single log line, used by
/// `squash_log` to rank which lines are most worth keeping.
#[derive(Clone, Copy, PartialEq, Eq)]
enum LogLevel {
    Error,
    Warn,
    Info,
    Debug,
}

/// Classify a single log line by scanning for level keywords.
///
/// Why substring matching on a lowercased copy instead of a real log
/// parser: tool output comes from arbitrary external processes with no
/// consistent log format, so keyword sniffing is the only detector that
/// generalizes across all of them.
fn classify_line(line: &str) -> LogLevel {
    let lower = line.to_lowercase();
    if lower.contains("error") || lower.contains("fail") || lower.contains("panic") {
        LogLevel::Error
    } else if lower.contains("warn") {
        LogLevel::Warn
    } else if lower.contains("debug") || lower.contains("trace") {
        LogLevel::Debug
    } else {
        LogLevel::Info
    }
}

/// Heuristic gate for routing to `squash_log` vs `squash_generic`: at
/// least 3 lines that look like error/warning/stack-trace output.
fn looks_log_shaped(text: &str) -> bool {
    let hits = text
        .lines()
        .filter(|l| {
            let lower = l.to_lowercase();
            lower.contains("error")
                || lower.contains("warn")
                || lower.contains("fail")
                || lower.contains("panic")
                || l.trim_start().starts_with("at ")
        })
        .count();
    hits >= 3
}

/// Compress log-shaped output: keep up to 20 highest-scored error lines
/// and up to 10 highest-scored warning lines (score = level weight +
/// 0.3 if the line looks like a stack-trace frame), each with a
/// +/-2-line context window, replacing every gap with a `[N lines
/// omitted]` marker.
///
/// Why not a comment-shaped marker (e.g. `// N lines omitted`): the
/// `rtk` project's own regression tests found that shape gets parsed by
/// the LLM as code and triggers a retry loop.
fn squash_log(text: &str) -> String {
    let lines: Vec<&str> = text.lines().collect();
    let levels: Vec<LogLevel> = lines.iter().map(|l| classify_line(l)).collect();

    let score = |i: usize| -> f32 {
        let level_score = match levels[i] {
            LogLevel::Error => 1.0,
            LogLevel::Warn => 0.5,
            LogLevel::Info => 0.1,
            LogLevel::Debug => 0.05,
        };
        let stack_boost = if lines[i].trim_start().starts_with("at ") {
            0.3
        } else {
            0.0
        };
        level_score + stack_boost
    };

    let mut error_idxs: Vec<usize> = (0..lines.len())
        .filter(|&i| levels[i] == LogLevel::Error)
        .collect();
    error_idxs.sort_by(|&a, &b| {
        score(b)
            .partial_cmp(&score(a))
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    error_idxs.truncate(20);

    let mut warn_idxs: Vec<usize> = (0..lines.len())
        .filter(|&i| levels[i] == LogLevel::Warn)
        .collect();
    warn_idxs.sort_by(|&a, &b| {
        score(b)
            .partial_cmp(&score(a))
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    warn_idxs.truncate(10);

    let mut keep: HashSet<usize> = HashSet::new();
    for &i in error_idxs.iter().chain(warn_idxs.iter()) {
        let lo = i.saturating_sub(2);
        let hi = (i + 2).min(lines.len().saturating_sub(1));
        keep.extend(lo..=hi);
    }

    if keep.is_empty() {
        return squash_generic(text, GENERIC_BUDGET_BYTES);
    }

    render_kept_lines(&lines, &keep)
}

/// Importance-ranked truncation for content that isn't JSON or
/// log-shaped: keep the first 10 and last 10 lines, plus any
/// non-blank line that isn't a repeat of the one before it, until
/// `budget` bytes are used.
fn squash_generic(text: &str, budget: usize) -> String {
    let lines: Vec<&str> = text.lines().collect();
    if lines.len() <= 20 {
        return text.chars().take(budget).collect();
    }

    let head_end = 10;
    let tail_start = lines.len() - 10;
    let mut keep: HashSet<usize> = (0..head_end).chain(tail_start..lines.len()).collect();

    let mut used: usize = lines[..head_end].iter().map(|l| l.len() + 1).sum::<usize>()
        + lines[tail_start..]
            .iter()
            .map(|l| l.len() + 1)
            .sum::<usize>();
    let mut prev = "";
    for (i, &line) in lines.iter().enumerate().take(tail_start).skip(head_end) {
        let non_trivial = !line.trim().is_empty() && line != prev;
        if non_trivial && used + line.len() < budget {
            keep.insert(i);
            used += line.len() + 1;
        }
        prev = line;
    }

    render_kept_lines(&lines, &keep)
}

/// Render a subset of `lines` in order, inserting a `[N lines omitted]`
/// marker at every gap between kept lines.
fn render_kept_lines(lines: &[&str], keep: &HashSet<usize>) -> String {
    let mut kept_sorted: Vec<usize> = keep.iter().copied().collect();
    kept_sorted.sort_unstable();

    let mut out = String::new();
    let mut cursor = 0usize;
    for &i in &kept_sorted {
        if i > cursor {
            let _ = writeln!(out, "[{} lines omitted]", i - cursor);
        }
        out.push_str(lines[i]);
        out.push('\n');
        cursor = i + 1;
    }
    if cursor < lines.len() {
        let _ = writeln!(out, "[{} lines omitted]", lines.len() - cursor);
    }
    out
}

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

    #[test]
    fn output_under_the_floor_passes_through_unchanged() {
        let small = "short output";
        assert_eq!(apply("bash", small), small);
    }

    #[test]
    fn read_tool_output_is_never_squashed_even_when_huge_json() {
        let big_json = format!(
            "{{\"description\": \"{}\"}}",
            "a very long description value that repeats ".repeat(100),
        );
        assert!(big_json.len() > SQUASH_FLOOR_BYTES);
        assert_eq!(apply("read", &big_json), big_json);
    }

    #[test]
    fn json_output_over_floor_keeps_structure_and_short_values() {
        let value = serde_json::json!({
            "id": "abc123",
            "note": "hi",
            "description": "a very long description value that repeats ".repeat(100),
        });
        let text = serde_json::to_string(&value).unwrap();
        assert!(text.len() > SQUASH_FLOOR_BYTES);

        let result = apply("some_mcp_tool", &text);
        let parsed: serde_json::Value =
            serde_json::from_str(&result).expect("squashed JSON must still be valid JSON");

        assert_eq!(parsed["id"], "abc123", "short values must survive");
        assert_eq!(parsed["note"], "hi", "short values must survive");
        assert_ne!(
            parsed["description"].as_str().unwrap().len(),
            value["description"].as_str().unwrap().len(),
            "long low-entropy value must be shrunk",
        );
    }

    #[test]
    fn json_array_elements_past_third_are_squashed_harder() {
        // A UUID-shaped value has no internal whitespace and clears the
        // entropy threshold, so under the *normal* per-value rule (which
        // still applies to array indices 0-2) it survives untouched.
        // Padding elsewhere in the object pushes total size over the
        // squash floor without affecting which array elements get kept.
        let identifier = "550e8400-e29b-41d4-a716-446655440000";
        let padding = "padding text to push this payload past the squash floor so apply() actually dispatches to squash_json ".repeat(20);
        let value = serde_json::json!({
            "padding": padding,
            "items": [identifier, identifier, identifier, identifier],
        });
        let text = serde_json::to_string(&value).unwrap();
        assert!(text.len() > SQUASH_FLOOR_BYTES);

        let result = apply("some_mcp_tool", &text);
        let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
        let items = parsed["items"].as_array().unwrap();

        assert_eq!(items[0].as_str().unwrap(), identifier, "index 0 is under the array cutoff and identifier-shaped, so it's kept under the normal rule");
        assert_eq!(
            items[2].as_str().unwrap(),
            identifier,
            "index 2 is still under the cutoff (past-third means index >= 3)"
        );
        assert_ne!(items[3].as_str().unwrap(), identifier, "index 3 must be force-elided even though it's identifier-shaped and would survive at any earlier index");
    }

    #[test]
    fn log_like_output_keeps_error_lines_and_marks_omissions() {
        // `looks_log_shaped` requires >= 3 lines matching error/warn/fail/
        // panic/stack-frame patterns before routing to `squash_log` at
        // all — a single error line isn't enough and would silently fall
        // through to `squash_generic` instead, so this fixture needs at
        // least 3 such lines, spread apart, to actually exercise
        // squash_log's scoring/windowing logic (not just its fallback).
        let mut lines = vec!["build started".to_string()];
        for i in 0..200 {
            lines.push(format!("info: compiling module {i}"));
        }
        lines.push("error: something failed early in the build".to_string());
        for i in 0..200 {
            lines.push(format!("info: compiling module {}", i + 200));
        }
        lines.push("warning: deprecated api used somewhere".to_string());
        lines.push("error: something failed at the end".to_string());
        let text = lines.join("\n");
        assert!(text.len() > SQUASH_FLOOR_BYTES);

        let result = apply("bash", &text);

        assert!(result.contains("error: something failed early in the build"));
        assert!(result.contains("error: something failed at the end"));
        assert!(result.contains("lines omitted"));
        assert!(result.len() < text.len());
    }

    #[test]
    fn non_bash_tool_with_log_shaped_content_is_not_log_compressed() {
        // A grep result whose matched lines all mention "error" would
        // trip `looks_log_shaped`'s >=3-line keyword threshold just like
        // a real build log — but `squash_log` caps at 20 highest-scored
        // error lines with no guaranteed tail retention, silently
        // dropping legitimate matches past that cap. Only `bash` is
        // treated as log-shaped; `grep` must fall through to
        // `squash_generic`, which always keeps the first and last 10
        // lines regardless of score. With every line tied at the same
        // score, a `squash_log` route would keep indices 0-19 (stable
        // sort preserves original order on ties) and drop index 49 —
        // so asserting the tail survives is a route-distinguishing
        // check, not just a content check.
        let lines: Vec<String> = (0..50)
            .map(|i| format!("src/file{i}.rs:{i}: error handling for case {i}"))
            .collect();
        let text = lines.join("\n");
        assert!(text.len() > SQUASH_FLOOR_BYTES);

        let result = apply("grep", &text);

        assert!(
            result.contains("src/file0.rs:0: error handling for case 0"),
            "generic keeps head"
        );
        assert!(
            result.contains("src/file49.rs:49: error handling for case 49"),
            "generic keeps tail — squash_log would have dropped this"
        );
    }

    #[test]
    fn generic_large_text_is_truncated_with_omission_marker() {
        let lines: Vec<String> = (0..500)
            .map(|i| format!("line number {i} of plain output"))
            .collect();
        let text = lines.join("\n");
        assert!(text.len() > SQUASH_FLOOR_BYTES);

        let result = apply("bash", &text);

        assert!(
            result.contains("line number 0 of plain output"),
            "keeps head"
        );
        assert!(
            result.contains("line number 499 of plain output"),
            "keeps tail"
        );
        assert!(result.contains("lines omitted"));
        assert!(result.len() < text.len());
    }
}