feat(context): tambah context::squash untuk kompresi hasil tool
Kompresi per-jenis-konten (JSON: pertahankan struktur & value pendek/ entropi tinggi, buang value panjang bertele-tele; log: simpan baris error/warning berskor tertinggi + konteks sekitarnya; generic: potong importance-ranked) untuk hasil tool di atas 1.5KB. Tool read dikecualikan total karena isinya harus tetap byte-exact untuk edit selanjutnya.
This commit is contained in:
asepharyana
@@ -11,5 +11,6 @@
|
||||
//! `TurnEvent::Compacted`.
|
||||
|
||||
pub mod dedup;
|
||||
pub mod squash;
|
||||
pub mod tokens;
|
||||
pub mod window;
|
||||
|
||||
@@ -0,0 +1,393 @@
|
||||
//! 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` : 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;
|
||||
|
||||
/// 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"];
|
||||
|
||||
/// 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 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() + 1 <= 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 {
|
||||
out.push_str(&format!("[{} lines omitted]\n", i - cursor));
|
||||
}
|
||||
out.push_str(lines[i]);
|
||||
out.push('\n');
|
||||
cursor = i + 1;
|
||||
}
|
||||
if cursor < lines.len() {
|
||||
out.push_str(&format!("[{} lines omitted]\n", 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 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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user