use crate::message::{Message, Role};
pub const MICROCOMPACT_PLACEHOLDER: &str = "[older tool output trimmed to fit budget]";
pub const MIN_PRUNE_LENGTH: usize = 200;
#[derive(Debug, Clone)]
pub struct Microcompactor {
pub trigger_tool_count: usize,
pub keep_recent: usize,
}
impl Default for Microcompactor {
fn default() -> Self {
Self {
trigger_tool_count: 12,
keep_recent: 4,
}
}
}
impl Microcompactor {
pub fn new(trigger_tool_count: usize, keep_recent: usize) -> Self {
Self {
trigger_tool_count,
keep_recent,
}
}
pub fn prune(&self, messages: &mut [Message]) -> usize {
let tool_indices: Vec<usize> = messages
.iter()
.enumerate()
.filter(|(_, m)| m.role == Role::Tool)
.map(|(i, _)| i)
.collect();
let count = tool_indices.len();
if count <= self.trigger_tool_count {
return 0;
}
let eligible_count = count.saturating_sub(self.keep_recent);
let target_prune = count.saturating_sub(self.trigger_tool_count);
let to_process = eligible_count.min(target_prune);
let mut pruned = 0;
for idx in tool_indices.iter().take(to_process) {
let idx = *idx;
let msg = &mut messages[idx];
if msg.content == MICROCOMPACT_PLACEHOLDER || msg.content.len() <= MIN_PRUNE_LENGTH {
continue;
}
msg.content = MICROCOMPACT_PLACEHOLDER.to_string();
pruned += 1;
}
pruned
}
}
pub fn build_microcompactor_from_env(
trigger_raw: Option<&str>,
keep_raw: Option<&str>,
) -> Option<Microcompactor> {
let trigger: usize = match trigger_raw {
None | Some("0") | Some("off") | Some("false") => return None, Some(s) => s.parse::<usize>().ok().filter(|&n| n > 0)?,
};
let keep: usize = match keep_raw {
Some(s) => s.parse::<usize>().ok().unwrap_or(4),
None => 4,
};
Some(Microcompactor::new(trigger, keep))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::message::Message;
fn tool_msg(content_len: usize) -> Message {
Message::tool_result("call_1", "x".repeat(content_len))
}
#[test]
fn prune_noop_when_under_trigger() {
let mut msgs: Vec<Message> = (0..5).map(|_| tool_msg(300)).collect();
let mc = Microcompactor::new(12, 4);
let pruned = mc.prune(&mut msgs);
assert_eq!(pruned, 0);
for msg in &msgs {
assert_ne!(msg.content, MICROCOMPACT_PLACEHOLDER);
}
}
#[test]
fn prune_oldest_keeps_recent() {
let mut msgs: Vec<Message> = (0..15).map(|_| tool_msg(300)).collect();
let mc = Microcompactor::new(12, 4);
let pruned = mc.prune(&mut msgs);
assert_eq!(
pruned, 3,
"should prune exactly 3 to get back under trigger"
);
for msg in msgs[11..15].iter() {
assert_ne!(
msg.content, MICROCOMPACT_PLACEHOLDER,
"message must not be pruned"
);
}
for msg in msgs[0..3].iter() {
assert_eq!(
msg.content, MICROCOMPACT_PLACEHOLDER,
"message must be pruned"
);
}
for msg in msgs[3..11].iter() {
assert_ne!(
msg.content, MICROCOMPACT_PLACEHOLDER,
"message should not be pruned (pruning stopped at 3)"
);
}
}
#[test]
fn prune_preserves_tool_messages() {
let mut msgs: Vec<Message> = (0..15).map(|_| tool_msg(300)).collect();
let total_before = msgs.len();
let mc = Microcompactor::new(12, 4);
let _pruned = mc.prune(&mut msgs);
assert_eq!(msgs.len(), total_before, "message count must not change");
for msg in &msgs {
assert_eq!(msg.role, Role::Tool, "all messages must still be Tool role");
}
}
#[test]
fn prune_skips_short_results() {
let mut msgs: Vec<Message> = Vec::new();
for _ in 0..3 {
msgs.push(tool_msg(50));
}
for _ in 0..8 {
msgs.push(tool_msg(300));
}
for _ in 0..4 {
msgs.push(tool_msg(300));
}
let mc = Microcompactor::new(12, 4);
let pruned = mc.prune(&mut msgs);
assert_eq!(pruned, 0, "short results should not be pruned");
}
#[test]
fn prune_skips_short_results_and_prunes_eligible_ones() {
let mut msgs: Vec<Message> = Vec::new();
msgs.push(tool_msg(50));
msgs.push(tool_msg(50));
for _ in 0..5 {
msgs.push(tool_msg(300));
}
for _ in 0..4 {
msgs.push(tool_msg(300));
}
let mc = Microcompactor::new(8, 4);
let pruned = mc.prune(&mut msgs);
assert_ne!(msgs[0].content, MICROCOMPACT_PLACEHOLDER);
assert_ne!(msgs[1].content, MICROCOMPACT_PLACEHOLDER);
assert_eq!(msgs[2].content, MICROCOMPACT_PLACEHOLDER);
for msg in msgs[7..11].iter() {
assert_ne!(msg.content, MICROCOMPACT_PLACEHOLDER);
}
assert!(pruned >= 1, "should prune at least 1 eligible long result");
}
#[test]
fn prune_idempotent() {
let mut msgs: Vec<Message> = (0..15).map(|_| tool_msg(300)).collect();
let mc = Microcompactor::new(12, 4);
let first = mc.prune(&mut msgs);
assert_eq!(first, 3, "first prune should prune 3");
let second = mc.prune(&mut msgs);
assert_eq!(second, 0, "second prune should be no-op (idempotent)");
}
#[test]
fn prune_exact_at_trigger_no_prune() {
let mut msgs: Vec<Message> = (0..12).map(|_| tool_msg(300)).collect();
let mc = Microcompactor::new(12, 4);
let pruned = mc.prune(&mut msgs);
assert_eq!(pruned, 0);
}
#[test]
fn prune_just_over_trigger_prunes_appropriate_number() {
let mut msgs: Vec<Message> = (0..13).map(|_| tool_msg(300)).collect();
let mc = Microcompactor::new(12, 4);
let pruned = mc.prune(&mut msgs);
assert_eq!(pruned, 1, "13 > 12 → should prune exactly 1");
}
#[test]
fn build_microcompactor_from_env_disabled_when_zero() {
assert!(build_microcompactor_from_env(Some("0"), None).is_none());
assert!(build_microcompactor_from_env(Some("off"), None).is_none());
assert!(build_microcompactor_from_env(Some("false"), None).is_none());
}
#[test]
fn build_microcompactor_from_env_disabled_by_default() {
assert!(
build_microcompactor_from_env(None, None).is_none(),
"unset trigger must return None (disabled by default)"
);
}
#[test]
fn build_microcompactor_from_env_explicit_values() {
let mc = build_microcompactor_from_env(Some("20"), Some("6")).unwrap();
assert_eq!(mc.trigger_tool_count, 20);
assert_eq!(mc.keep_recent, 6);
}
#[test]
fn build_microcompactor_from_env_trigger_only() {
let mc = build_microcompactor_from_env(Some("15"), None).unwrap();
assert_eq!(mc.trigger_tool_count, 15);
assert_eq!(mc.keep_recent, 4);
}
#[test]
fn build_microcompactor_from_env_invalid_trigger_returns_none() {
assert!(
build_microcompactor_from_env(Some("not-a-number"), None).is_none(),
"unparseable trigger must return None"
);
}
}