From ea0b4988299894e4588d851ebc91704a9e73bc73 Mon Sep 17 00:00:00 2001 From: asepharyana Date: Thu, 16 Jul 2026 04:42:35 +0700 Subject: [PATCH] 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. --- src/app/runtime/context/mod.rs | 1 + src/app/runtime/context/squash.rs | 393 ++++++++++++++++++++++++++++++ 2 files changed, 394 insertions(+) create mode 100644 src/app/runtime/context/squash.rs diff --git a/src/app/runtime/context/mod.rs b/src/app/runtime/context/mod.rs index b0e5c74..a0181f6 100644 --- a/src/app/runtime/context/mod.rs +++ b/src/app/runtime/context/mod.rs @@ -11,5 +11,6 @@ //! `TurnEvent::Compacted`. pub mod dedup; +pub mod squash; pub mod tokens; pub mod window; diff --git a/src/app/runtime/context/squash.rs b/src/app/runtime/context/squash.rs new file mode 100644 index 0000000..06a4930 --- /dev/null +++ b/src/app/runtime/context/squash.rs @@ -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::(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::(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 = 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 = 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 = (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 = (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 = 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 = (0..head_end).chain(tail_start..lines.len()).collect(); + + let mut used: usize = lines[..head_end].iter().map(|l| l.len() + 1).sum::() + + lines[tail_start..].iter().map(|l| l.len() + 1).sum::(); + 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) -> String { + let mut kept_sorted: Vec = 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 = (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()); + } +}